From e3e717b948074b0bca9629f936a1e199b86ed94b Mon Sep 17 00:00:00 2001 From: Conatum Date: Fri, 27 Oct 2023 11:20:05 +0300 Subject: [PATCH] (core): Add explicit cog type overloads. --- src/core/cog.py | 1 - src/meta/LionBot.py | 138 +++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 135 insertions(+), 4 deletions(-) diff --git a/src/core/cog.py b/src/core/cog.py index 1f36ffb3..dde79881 100644 --- a/src/core/cog.py +++ b/src/core/cog.py @@ -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 diff --git a/src/meta/LionBot.py b/src/meta/LionBot.py index bb48bcf9..9d514c67 100644 --- a/src/meta/LionBot.py +++ b/src/meta/LionBot.py @@ -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__}")