import logging from dataclasses import dataclass from typing import Union from bson import ObjectId from modules.database import col_users logger = logging.getLogger(__name__) @dataclass class PyroUser: """Dataclass of DB entry of a user""" __slots__ = ("_id", "id", "locale") _id: ObjectId id: int locale: Union[str, None] @classmethod async def find(cls, id: int, locale: Union[str, None] = None): """Find user in database and create new record if user does not exist. ### Args: * id (`int`): User's Telegram ID * locale (`Union[str, None]`, *optional*): User's locale. Defaults to `None`. ### Raises: * `RuntimeError`: Raised when user entry after insertion could not be found. ### Returns: * `PyroUser`: User with its database data. """ db_entry = await col_users.find_one({"id": id}) if db_entry is None: inserted = await col_users.insert_one({"id": id, "locale": locale}) db_entry = await col_users.find_one({"_id": inserted.inserted_id}) if db_entry is None: raise RuntimeError("Could not find inserted user entry.") return cls(**db_entry) async def update_locale(self, locale: Union[str, None]) -> None: """Change user's locale stored in the database. ### Args: * locale (`Union[str, None]`): New locale to be set. """ logger.debug("%s's locale has been set to %s", self.id, locale) await col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}})