PyrogramBotBase/classes/pyrouser_mongo.py

import logging
from dataclasses import dataclass
from typing import Union

from bson import ObjectId

from modules.database_mongo 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}})