(core): Add explicit cog type overloads.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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__}")
|
||||
|
||||
Reference in New Issue
Block a user