TelegramBot/classes/pyrouser.py

import logging
from dataclasses import dataclass
from typing import Any, Union

from bson import ObjectId

from classes.location import Location
from modules.database import col_users

logger = logging.getLogger(__name__)


@dataclass
class PyroUser:
    """Dataclass of DB entry of a user"""

    __slots__ = (
        "_id",
        "id",
        "locale",
        "enabled",
        "location",
        "offset",
        "time_hour",
        "time_minute",
    )

    _id: ObjectId
    id: int
    locale: Union[str, None]
    enabled: bool
    location: Union[Location, None]
    offset: int
    time_hour: int
    time_minute: int

    @classmethod
    async def find(
        cls,
        id: int,
        locale: Union[str, None] = None,
        enabled: bool = True,
        location_id: int = 0,
        offset: int = 1,
        time_hour: int = 18,
        time_minute: int = 0,
    ):
        db_entry = await col_users.find_one({"id": id})

        if db_entry is None:
            inserted = await col_users.insert_one(
                {
                    "id": id,
                    "locale": locale,
                    "enabled": enabled,
                    "location": location_id,
                    "offset": offset,
                    "time_hour": time_hour,
                    "time_minute": time_minute,
                }
            )
            db_entry = await col_users.find_one({"_id": inserted.inserted_id})

        if db_entry is None:
            raise RuntimeError("Could not find inserted user entry.")

        try:
            db_entry["location"] = await Location.get(db_entry["location"])  # type: ignore
        except ValueError:
            db_entry["location"] = None  # type: ignore

        return cls(**db_entry)

    @classmethod
    async def from_dict(cls, **kwargs):
        if "location" in kwargs:
            try:
                kwargs["location"] = await Location.get(kwargs["location"])  # type: ignore
            except ValueError:
                kwargs["location"] = None  # type: ignore
        return cls(**kwargs)

    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}})

    async def update_state(self, enabled: bool = False) -> None:
        logger.debug("%s's state has been set to %s", self.id, enabled)
        await col_users.update_one({"_id": self._id}, {"$set": {"enabled": enabled}})

    async def update_location(self, location_id: int = 0) -> None:
        logger.debug("%s's location has been set to %s", self.id, location_id)
        await col_users.update_one(
            {"_id": self._id}, {"$set": {"location": location_id}}
        )

    async def update_offset(self, offset: int = 1) -> None:
        logger.debug("%s's offset has been set to %s", self.id, offset)
        await col_users.update_one({"_id": self._id}, {"$set": {"offset": offset}})

    async def update_time(self, hour: int = 18, minute: int = 0) -> None:
        logger.debug("%s's time has been set to %s h. %s m.", self.id, hour, minute)
        await col_users.update_one(
            {"_id": self._id}, {"$set": {"time_hour": hour, "time_minute": minute}}
        )

    async def delete(self) -> None:
        logger.debug("%s's data has been deleted", self.id)
        await col_users.delete_one({"_id": self._id})

    async def checkout(self) -> Any:
        logger.debug("%s's data has been checked out", self.id)
        db_entry = await col_users.find_one({"_id": self._id})

        if db_entry is None:
            raise KeyError(
                f"DB record with id {self._id} of user {self.id} is not found"
            )

        del db_entry["_id"]  # type: ignore

        return db_entry