ShopBot/utils/database/user.py

from aiogram.types import User as UserType
from peewee import Model

from .model import User as UserModel
from .model import Admin, Operator
from utils import types


class Register:

    @classmethod
    def __register_user(cls, model: Model, user: UserType):
        model.insert(
            user_id=user.id,
            first_name=user.first_name,
            last_name=user.last_name,
            username=user.username
        ).on_conflict_replace().execute()
    
    @classmethod
    def register_user(cls, user: UserType):
        return cls.__register_user(UserModel, user)
    
    @classmethod
    def register_admin(cls, user: UserType):
        return cls.__register_user(Admin, user)

    @classmethod
    def register_operator(cls, user: UserType):
        return cls.__register_user(Operator, user)


class User:

    @staticmethod
    def get_user(user_id: int):
        if not UserModel.select().where(UserModel.user_id==user_id).exists():
            return None
        data = UserModel.get(UserModel.user_id==user_id)
        return types.User(
            user_id=data.user_id,
            first_name=data.first_name,
            last_name=data.last_name,
            username=data.username
        )
    
    @staticmethod
    def get_state(user_id:int):
        if Operator.select().where(Operator.user_id == user_id).exists():
            return Operator.get(Operator.user_id == user_id).active
        return False
    
    @staticmethod
    def set_state(user_id:int, state:bool=False):
        if Operator.select().where(Operator.user_id == user_id).exists():
            Operator.update(active=state).where(Operator.user_id == user_id).execute()
        return None