103 lines
2.9 KiB
Python
103 lines
2.9 KiB
Python
|
# -*- 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
|