# -*- coding: utf-8 -*- ########################### # CacheDB implementations # ########################### # Imports from tubot.db.abc import CacheDB from tubot.db.types import CacheDBTypes, UserStates, User from tubot.static.env import PN_CACHE from pickle import loads, dumps from aiofiles.ospath import isdir, isfile from aiofiles.os import mkdir from aiofiles import open from asyncio import sleep class PythonCache(CacheDB): """ Native python implementation of Cache DataBase """ CACHE_DIR = PN_CACHE()() _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"{self.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(self.CACHE_DIR): await mkdir(self.CACHE_DIR) async with open(f"{self.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(self.CACHE_DIR): await mkdir(self.CACHE_DIR) async with open(f"{self.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