TelegramPoster/classes/pyrouser.py

from dataclasses import dataclass
from datetime import datetime
from typing import Union

from bson import ObjectId
from libbot import config_get
from libbot.pyrogram.classes import PyroClient

from modules.database import col_users


@dataclass
class PyroUser:
    """Dataclass of DB entry of a user"""

    _id: ObjectId
    id: int
    locale: Union[str, None]
    banned: bool
    cooldown: datetime
    subscription: dict

    async def update_locale(self, locale: str) -> None:
        await col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}})

    async def update_cooldown(self, time: datetime = datetime.now()) -> None:
        await col_users.update_one({"_id": self._id}, {"$set": {"cooldown": time}})

    async def block(self) -> None:
        """Ban user from using command and submitting content."""
        await col_users.update_one({"_id": self._id}, {"$set": {"banned": True}})

    async def unblock(self) -> None:
        """Allow user to use command and submit posts again."""
        await col_users.update_one({"_id": self._id}, {"$set": {"banned": False}})

    async def is_limited(self, app: Union[PyroClient, None] = None) -> bool:
        """Check if user is on a cooldown after submitting something.

        ### Returns:
            `bool`: Must be `True` if on the cooldown and `False` if not
        """
        admins = (
            await config_get("admins", "bot") + [await config_get("owner", "bot")]
            if app is None
            else app.admins
        )

        return (datetime.now() - self.cooldown).total_seconds() < (
            app.config["submission"]["timeout"]
            if app is not None
            else await config_get("timeout", "submission")
        )