TelegramBot/classes/pyrouser.py

119 lines
3.6 KiB
Python

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