TelegramBot/classes/pyroclient.py

from typing import List, Union

from apscheduler.triggers.cron import CronTrigger
from libbot.pyrogram.classes import PyroClient as LibPyroClient
from pymongo import ASCENDING, GEOSPHERE, TEXT
from pyrogram.types import User

from classes.location import Location
from classes.pyrouser import PyroUser
from modules.database import col_locations
from modules.reminder import remind


class PyroClient(LibPyroClient):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        if self.scheduler is not None:
            self.scheduler.add_job(
                remind, CronTrigger.from_crontab("* * * * *"), args=(self,)
            )
        self.contexts = []

    async def start(self, **kwargs):
        await col_locations.create_index(
            [("id", ASCENDING)], name="location_id", unique=True
        )
        await col_locations.create_index(
            [("location", GEOSPHERE)],
            name="location_location",
        )
        await col_locations.create_index([("name", TEXT)], name="location_name")
        return await super().start(**kwargs)

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

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

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

        return (
            await PyroUser.find(user)
            if isinstance(user, int)
            else await PyroUser.find(user.id, locale=user.language_code)
        )

    async def get_location(self, id: int) -> Location:
        """Get Location by it's ID.

        ### Args:
            * id (`int`): Location's ID. Defaults to `None`.

        ### Returns:
            * `Location`: Location from database as an object.
        """

        return await Location.get(id)

    async def list_locations(self) -> List[Location]:
        """Get all locations stored in database.

        ### Returns:
            * `List[Location]`: List of `Location` objects.
        """
        return [
            await Location.get(record["id"]) async for record in col_locations.find({})
        ]