(core): Add explicit cog type overloads.

This commit is contained in:
2023-10-27 11:20:05 +03:00
parent bbd00267fa
commit e3e717b948
2 changed files with 135 additions and 4 deletions

View File

@@ -74,7 +74,6 @@ class CoreCog(LionCog):
self.bot.add_listener(self.shard_update_guilds, name='on_guild_join')
self.bot.add_listener(self.shard_update_guilds, name='on_guild_remove')
self.bot.core = self
await self.bot.add_cog(self.lions)
# Load the app command cache

View File

@@ -1,4 +1,4 @@
from typing import List, Optional, TYPE_CHECKING
from typing import List, Literal, LiteralString, Optional, TYPE_CHECKING, overload
import logging
import asyncio
from weakref import WeakValueDictionary
@@ -24,7 +24,30 @@ from .errors import HandledException, SafeCancellation
from .monitor import SystemMonitor, ComponentMonitor, StatusLevel, ComponentStatus
if TYPE_CHECKING:
from core import CoreCog
from core.cog import CoreCog
from core.config import ConfigCog
from tracking.voice.cog import VoiceTrackerCog
from tracking.text.cog import TextTrackerCog
from modules.config.cog import GuildConfigCog
from modules.economy.cog import Economy
from modules.member_admin.cog import MemberAdminCog
from modules.meta.cog import MetaCog
from modules.moderation.cog import ModerationCog
from modules.pomodoro.cog import TimerCog
from modules.premium.cog import PremiumCog
from modules.ranks.cog import RankCog
from modules.reminders.cog import Reminders
from modules.rooms.cog import RoomCog
from modules.schedule.cog import ScheduleCog
from modules.shop.cog import ShopCog
from modules.skins.cog import CustomSkinCog
from modules.sponsors.cog import SponsorCog
from modules.statistics.cog import StatsCog
from modules.sysadmin.dash import LeoSettings
from modules.tasklist.cog import TasklistCog
from modules.topgg.cog import TopggCog
from modules.user_config.cog import UserConfigCog
from modules.video_channels.cog import VideoCog
logger = logging.getLogger(__name__)
@@ -46,7 +69,6 @@ class LionBot(Bot):
# self.appdata = appdata
self.config = config
self.app_ipc = app_ipc
self.core: 'CoreCog' = None
self.translator = translator
self.system_monitor = SystemMonitor()
@@ -58,6 +80,10 @@ class LionBot(Bot):
self._talk_global_dispatch = app_ipc.register_route('dispatch')(self._handle_global_dispatch)
@property
def core(self):
return self.get_cog('CoreCog')
async def _handle_global_dispatch(self, event_name: str, *args, **kwargs):
self.dispatch(event_name, *args, **kwargs)
@@ -107,6 +133,112 @@ class LionBot(Bot):
self.tree.copy_global_to(guild=guild)
await self.tree.sync(guild=guild)
# To make the type checker happy about fetching cogs by name
# TODO: Move this to stubs at some point
@overload
def get_cog(self, name: Literal['CoreCog']) -> 'CoreCog':
...
@overload
def get_cog(self, name: Literal['ConfigCog']) -> 'ConfigCog':
...
@overload
def get_cog(self, name: Literal['VoiceTrackerCog']) -> 'VoiceTrackerCog':
...
@overload
def get_cog(self, name: Literal['TextTrackerCog']) -> 'TextTrackerCog':
...
@overload
def get_cog(self, name: Literal['GuildConfigCog']) -> 'GuildConfigCog':
...
@overload
def get_cog(self, name: Literal['Economy']) -> 'Economy':
...
@overload
def get_cog(self, name: Literal['MemberAdminCog']) -> 'MemberAdminCog':
...
@overload
def get_cog(self, name: Literal['MetaCog']) -> 'MetaCog':
...
@overload
def get_cog(self, name: Literal['ModerationCog']) -> 'ModerationCog':
...
@overload
def get_cog(self, name: Literal['TimerCog']) -> 'TimerCog':
...
@overload
def get_cog(self, name: Literal['PremiumCog']) -> 'PremiumCog':
...
@overload
def get_cog(self, name: Literal['RankCog']) -> 'RankCog':
...
@overload
def get_cog(self, name: Literal['Reminders']) -> 'Reminders':
...
@overload
def get_cog(self, name: Literal['RoomCog']) -> 'RoomCog':
...
@overload
def get_cog(self, name: Literal['ScheduleCog']) -> 'ScheduleCog':
...
@overload
def get_cog(self, name: Literal['ShopCog']) -> 'ShopCog':
...
@overload
def get_cog(self, name: Literal['CustomSkinCog']) -> 'CustomSkinCog':
...
@overload
def get_cog(self, name: Literal['SponsorCog']) -> 'SponsorCog':
...
@overload
def get_cog(self, name: Literal['StatsCog']) -> 'StatsCog':
...
@overload
def get_cog(self, name: Literal['LeoSettings']) -> 'LeoSettings':
...
@overload
def get_cog(self, name: Literal['TasklistCog']) -> 'TasklistCog':
...
@overload
def get_cog(self, name: Literal['TopggCog']) -> 'TopggCog':
...
@overload
def get_cog(self, name: Literal['UserConfigCog']) -> 'UserConfigCog':
...
@overload
def get_cog(self, name: Literal['VideoCog']) -> 'VideoCog':
...
@overload
def get_cog(self, name: str) -> Optional[Cog]:
...
def get_cog(self, name: str) -> Optional[Cog]:
return super().get_cog(name)
async def add_cog(self, cog: Cog, **kwargs):
sup = super()
@log_wrap(action=f"Attach {cog.__cog_name__}")