Javelina/classes/pycordbot.py

from typing import Any, Union

from aiohttp import ClientSession
from discord import User
from libbot.pycord.classes import PycordBot as LibPycordBot

from classes.pycorduser import PycordUser
from modules.tracking.dhl import update_tracks_dhl


class PycordBot(LibPycordBot):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.client_session = ClientSession()

        if self.scheduler is None:
            return

        # Scheduler job for DHL parcel tracking
        self.scheduler.add_job(
            update_tracks_dhl,
            trigger="cron",
            hour=self.config["modules"]["tracking"]["fetch_hours"],
            args=[self, self.client_session],
        )

    async def find_user(self, user: Union[int, User]) -> PycordUser:
        """Find User by it's ID or User object.

        ### Args:
            * user (`Union[int, User]`): ID or User object to extract ID from.

        ### Returns:
            * `PycordUser`: User in database representation.
        """

        return (
            await PycordUser.find(user)
            if isinstance(user, int)
            else await PycordUser.find(user.id)
        )

    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)