Rewrite CacheDB && Implements UserStates

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
 

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
 

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)