from typing import Any, Union from aiohttp import ClientSession from bson import ObjectId from libbot.pycord.classes import PycordBot as LibPycordBot from classes.enums import MemberStatus from classes.errors import GuildNotFoundError, MemberNotFoundError from classes.pycord_guild import PycordGuild from classes.pycord_member import PycordMember class PycordBot(LibPycordBot): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.client_session = ClientSession() if self.scheduler is None: return async def find_guild( self, id: Union[int, ObjectId], create: bool = False, ) -> PycordGuild: """Find guild by their ID. ### Args: * id (`Union[int, ObjectId]`): Guild's ID * create (`bool`): Create DB entry if does not exist. Defaults to `False`. ### Raises: * `GuildNotFoundError`: Raised when guild entry could not be found. ### Returns: * `PycordGuild`: Guild in database representation. """ try: guild = await PycordGuild.find(id) except GuildNotFoundError as exc: if create and isinstance(id, int): return await PycordGuild.create(id) else: raise exc from exc return guild async def find_member( self, id: Union[int, ObjectId], guild: ObjectId, create: bool = False, ) -> PycordMember: """Find member by their ID and guild. ### Args: * id (`Union[int, ObjectId]`): Member's ID * guild (`ObjectId`): Discord guild's database ID * create (`bool`): Create DB entry if does not exist. Defaults to `False`. ### Raises: * `MemberNotFoundError`: Raised when member entry could not be found. ### Returns: * `PycordMember`: Member in database representation. """ try: member = await PycordMember.find(id, guild) except MemberNotFoundError as exc: if create and isinstance(id, int): return await PycordMember.create(id, guild, MemberStatus.UNVERIFIED) else: raise exc from exc return member async def close(self, *args: Any, **kwargs: Any) -> None: await self.client_session.close() if self.scheduler is not None: self.scheduler.shutdown() await super().close(*args, **kwargs)