EmojiCaptchaBot/classes/pyrouser.py

import logging
from dataclasses import dataclass
from typing import List

from bson import ObjectId

from modules.database import col_users

logger = logging.getLogger(__name__)


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

    __slots__ = (
        "_id",
        "id",
        "group",
        "failed",
        "emojis",
        "selected",
        "score",
        "mistakes",
    )

    _id: ObjectId
    id: int
    group: int
    failed: int
    emojis: List[str]
    selected: List[str]
    score: int
    mistakes: int

    @classmethod
    async def create_if_not_exists(
        cls,
        id: int,
        group: int,
        failed: bool = False,
        emojis: List[str] = [],
        selected: List[str] = [],
        score: int = 0,
        mistakes: int = 0,
    ):
        db_entry = await col_users.find_one({"id": id, "group": group})

        if db_entry is None:
            inserted = await col_users.insert_one(
                {
                    "id": id,
                    "group": group,
                    "failed": failed,
                    "emojis": emojis,
                    "selected": selected,
                    "score": score,
                    "mistakes": mistakes,
                }
            )
            db_entry = await col_users.find_one({"_id": inserted.inserted_id})

        if db_entry is None:
            raise RuntimeError("Could not find inserted user entry.")

        return cls(**db_entry)

    async def set_failed(self, failed: bool = True) -> None:
        logger.debug("%s's failure state has been set to %s", self.id, failed)
        await col_users.update_one({"_id": self._id}, {"$set": {"failed": failed}})

    async def set_emojis(self, emojis: List[str]) -> None:
        logger.debug("%s's emojis have been set to %s", self.id, emojis)
        await col_users.update_one({"_id": self._id}, {"$set": {"emojis": emojis}})

    async def set_score(self, score: int = 0) -> None:
        logger.debug("%s's score has been set to %s", self.id, score)
        await col_users.update_one({"_id": self._id}, {"$set": {"score": score}})

    async def set_mistakes(self, mistakes: int = 0) -> None:
        logger.debug("%s's mistakes count has been set to %s", self.id, mistakes)
        await col_users.update_one({"_id": self._id}, {"$set": {"mistakes": mistakes}})

    async def set_selected(self, emojis: List[str] = None) -> None:
        logger.debug("%s's elected emojis have been set to %s", self.id, emojis)
        await col_users.update_one({"_id": self._id}, {"$set": {"selected": emojis}})

    async def update_score(self, points: int = 1) -> None:
        logger.debug("%s point(s) have been added to %s score", points, self.id)
        await col_users.update_one(
            {"_id": self._id}, {"$set": {"score": self.score + points}}
        )

    async def update_mistakes(self, points: int = 1) -> None:
        logger.debug("%s point(s) have been added to %s mistakes", points, self.id)
        await col_users.update_one(
            {"_id": self._id}, {"$set": {"mistakes": self.mistakes + points}}
        )

    async def update_selected(self, entry: str) -> None:
        logger.debug("Emoji %s has been added to %s's selection list", entry, self.id)
        await col_users.update_one(
            {"_id": self._id}, {"$set": {"selected": self.selected + [entry]}}
        )