Add PythonCache implementation of CacheDB
- add User class to db.types - init db.cache - Add PythonCache class
This commit is contained in:
parent
ba696f4c56
commit
015e01b6be
|
@ -0,0 +1,102 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
|
||||
###########################
|
||||
# CacheDB implementations #
|
||||
###########################
|
||||
|
||||
# Imports
|
||||
from tubot.db.abc import CacheDB
|
||||
from tubot.db.types import CacheDBTypes, UserStates, User
|
||||
from pickle import loads, dumps
|
||||
from aiofiles.ospath import isdir, isfile
|
||||
from aiofiles.os import mkdir
|
||||
from aiofiles import open
|
||||
from asyncio import sleep
|
||||
|
||||
CACHE_DIR = "foo" # УДАЛИ ПОТОМ ЕБЛАН
|
||||
|
||||
|
||||
class PythonCache(CacheDB):
|
||||
"""
|
||||
Native python implementation of Cache DataBase
|
||||
"""
|
||||
|
||||
_ctype = CacheDBTypes.PythonPKL
|
||||
users: dict
|
||||
dirs: dict
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
async def _init(self) -> bool:
|
||||
self.users = {}
|
||||
self.dirs = {}
|
||||
if await isfile(f"{CACHE_DIR}/user_cache.pkl"):
|
||||
try:
|
||||
await self._load_pkl()
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _load_pkl(self) -> None:
|
||||
if not await isdir(CACHE_DIR):
|
||||
await mkdir(CACHE_DIR)
|
||||
async with open(f"{CACHE_DIR}/user_cache.pkl", "rb") as file:
|
||||
buffer = await file.read()
|
||||
pkl = loads(buffer)
|
||||
self.users = pkl
|
||||
|
||||
async def _save_pkl(self) -> None:
|
||||
if not await isdir(CACHE_DIR):
|
||||
await mkdir(CACHE_DIR)
|
||||
async with open(f"{CACHE_DIR}/user_cache.pkl", "wb") as file:
|
||||
await file.write(dumps(self.users))
|
||||
|
||||
async def __validate__(self) -> bool:
|
||||
return await self._init()
|
||||
|
||||
# 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
|
||||
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"]
|
||||
|
||||
# Dirs
|
||||
|
||||
async def cache_dirs(self, dirs: dict, expire: int) -> None:
|
||||
self.dirs = dirs
|
||||
await sleep(expire)
|
||||
self.dirs = {}
|
||||
|
||||
@property
|
||||
async def get_dirs(self) -> dict:
|
||||
return self.dirs
|
|
@ -12,11 +12,41 @@ class CacheDBTypes(Enum):
|
|||
"""
|
||||
Types of CacheDB
|
||||
"""
|
||||
...
|
||||
|
||||
PythonPKL = "python"
|
||||
|
||||
|
||||
class UserStates(Enum):
|
||||
"""
|
||||
Types of User status
|
||||
"""
|
||||
|
||||
...
|
||||
|
||||
|
||||
class User(object):
|
||||
"""
|
||||
User class
|
||||
"""
|
||||
|
||||
tg_id: int
|
||||
name: str
|
||||
state: UserStates | None = None
|
||||
auth: bool = False
|
||||
|
||||
def __init__(
|
||||
self, tg_id: int, name: str, state: UserStates | None = None, auth: bool = False
|
||||
) -> None:
|
||||
self.tg_id = tg_id
|
||||
self.name = name
|
||||
self.state = state
|
||||
self.auth = auth
|
||||
|
||||
@property
|
||||
def dict(self):
|
||||
return {
|
||||
"tg_id": self.tg_id,
|
||||
"name": self.name,
|
||||
"state": self.state,
|
||||
"auth": self.auth,
|
||||
}
|
||||
|
|
Reference in New Issue