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({}) ]