Rewrite CacheDB && Implements UserStates

This commit is contained in:
trueold89 2024-08-06 13:58:19 +03:00
parent dd27264fd5
commit 697dac6922
Signed by: trueold89
GPG Key ID: C122E85DD49E6B30
3 changed files with 37 additions and 104 deletions

View File

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

View File

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

View File

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