From 697dac692268d23279b358aa5e3ba38caa7d189b Mon Sep 17 00:00:00 2001 From: trueold89 Date: Tue, 6 Aug 2024 13:58:19 +0300 Subject: [PATCH] Rewrite CacheDB && Implements UserStates --- tubot/db/abc.py | 42 +++++--------------------- tubot/db/cache.py | 75 +++++++---------------------------------------- tubot/db/types.py | 24 +++++++++++---- 3 files changed, 37 insertions(+), 104 deletions(-) diff --git a/tubot/db/abc.py b/tubot/db/abc.py index 783d78c..d6bdc7a 100644 --- a/tubot/db/abc.py +++ b/tubot/db/abc.py @@ -7,7 +7,7 @@ # Imports from abc import ABC, abstractmethod from tubot.static.abc import IValidatable -from tubot.db.types import CacheDBTypes, UserStates +from tubot.db.types import CacheDBTypes, User class CacheDB(IValidatable, ABC): @@ -24,50 +24,22 @@ class CacheDB(IValidatable, ABC): # Users @abstractmethod - async def add_user(self, tg_id: int, name: str) -> None: + async def write_user(self, tg_id: int, user: User) -> None: """ - Add user to cache db + Writes user to cache db :param tg_id: User telegram id - :param name: User telegram name + :param user: User object """ raise NotImplementedError @abstractmethod - async def user_state(self, tg_id: int) -> UserStates: + async def read_user(self, tg_id: int) -> User: """ - Get user state - - :param tg_id: User telegram id - :return: Status of user - """ - raise NotImplementedError - - @abstractmethod - async def change_user_state(self, tg_id: int, status: UserStates) -> None: - """ - Change user state - - :param tg_id: User telegram id - :param status: New user status - """ - raise NotImplementedError - - @abstractmethod - async def auth_user(self, tg_id: int) -> None: - """ - Auth user - - :param tg_id: User telegram id - """ - raise NotImplementedError - - @abstractmethod - async def is_user_auth(self, tg_id: int) -> bool: - """ - Check if user is already auth + Writes user to cache db :param tg_id: User telegram id + :return: User object """ raise NotImplementedError diff --git a/tubot/db/cache.py b/tubot/db/cache.py index 3553ec9..d34c84c 100644 --- a/tubot/db/cache.py +++ b/tubot/db/cache.py @@ -60,38 +60,13 @@ class PythonCache(CacheDB): # Users - async def add_user(self, tg_id: int, name: str) -> None: - if tg_id in tuple(self.users.keys()): - raise ValueError("User already exists") - user = User(tg_id, name) - self.users[tg_id] = user.dict + async def write_user(self, tg_id: int, user: User) -> None: + self.users[tg_id] = user.to_dict await self._save_pkl() - async def user_state(self, tg_id: int) -> UserStates: - if tg_id not in tuple(self.users.keys()): - raise ValueError("User doesn't exists") - user = self.users[tg_id] - return user["state"] - - async def change_user_state(self, tg_id: int, status: UserStates) -> None: - if tg_id not in tuple(self.users.keys()): - raise ValueError("User doesn't exists") - user = self.users[tg_id] - user["state"] = status - await self._save_pkl() - - async def auth_user(self, tg_id: int) -> None: - if tg_id not in tuple(self.users.keys()): - raise ValueError("User doesn't exists") - user = self.users[tg_id] - user["auth"] = True - await self._save_pkl() - - async def is_user_auth(self, tg_id: int) -> bool: - if tg_id not in tuple(self.users.keys()): - raise ValueError("User doesn't exists") - user = self.users[tg_id] - return user["auth"] + async def read_user(self, tg_id: int) -> User: + user_data = self.users[tg_id] + return User.from_dict(user_data) # Dirs @@ -125,49 +100,21 @@ class RedisCache(CacheDB): # Users - async def add_user(self, tg_id: int, name: str) -> None: + async def write_user(self, tg_id: int, user: User) -> None: async with aioredis.from_url( f"redis://{self.host}", encoding="utf-8", decode_responses=True ) as redis: - user = User(tg_id, name) - json = json_dumps(user.dict) + json = json_dumps(user.to_dict) await redis.set(str(tg_id), json) - async def user_state(self, tg_id: int) -> UserStates: + async def read_user(self, tg_id: int) -> User: async with aioredis.from_url( f"redis://{self.host}", encoding="utf-8", decode_responses=True ) as redis: json = await redis.get(str(tg_id)) - user = json_loads(json) - return user["state"] - - async def change_user_state(self, tg_id: int, status: UserStates) -> None: - async with aioredis.from_url( - f"redis://{self.host}", encoding="utf-8", decode_responses=True - ) as redis: - json = await redis.get(str(tg_id)) - user = json_loads(json) - user["state"] = status - json = json_dumps(user) - await redis.set(str(tg_id), json) - - async def auth_user(self, tg_id: int) -> None: - async with aioredis.from_url( - f"redis://{self.host}", encoding="utf-8", decode_responses=True - ) as redis: - json = await redis.get(str(tg_id)) - user = json_loads(json) - user["auth"] = True - json = json_dumps(user) - await redis.set(str(tg_id), json) - - async def is_user_auth(self, tg_id: int) -> bool: - async with aioredis.from_url( - f"redis://{self.host}", encoding="utf-8", decode_responses=True - ) as redis: - json = await redis.get(str(tg_id)) - user = json_loads(json) - return bool(user["auth"]) + user_data = json_loads(json) + user_data["state"] = UserStates(user_data["state"]) + return User.from_dict(user_data) # Dirs diff --git a/tubot/db/types.py b/tubot/db/types.py index c17f6bc..75241d6 100644 --- a/tubot/db/types.py +++ b/tubot/db/types.py @@ -22,7 +22,9 @@ class UserStates(Enum): Types of User status """ - ... + IDLE = "IDLE" + DIRS = "DIRS" + WAIT_FOR_TORRENT = "WAIT_FOR_TORRENT" class User(object): @@ -32,11 +34,15 @@ class User(object): tg_id: int name: str - state: UserStates | None = None + state: UserStates = UserStates.IDLE auth: bool = False def __init__( - self, tg_id: int, name: str, state: UserStates | None = None, auth: bool = False + self, + tg_id: int, + name: str, + state: UserStates = UserStates.IDLE, + auth: bool = False, ) -> None: self.tg_id = tg_id self.name = name @@ -44,10 +50,18 @@ class User(object): self.auth = auth @property - def dict(self): + def to_dict(self): return { "tg_id": self.tg_id, "name": self.name, - "state": self.state, + "state": self.state.value, "auth": self.auth, } + + @classmethod + def from_dict(cls, usr: dict) -> "User": + tg = usr["tg_id"] + name = usr["name"] + state = UserStates(usr["state"]) + auth = usr["auth"] + return cls(tg, name, state, auth)