Initial commit

classes/captcha.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+from io import BytesIO
+from typing import List
+
+
+@dataclass
+class Captcha:
+    __slots__ = ("image", "emojis_all", "emojis_correct")
+    image: BytesIO
+    emojis_all: List[str]
+    emojis_correct: List[str]

classes/pyroclient.py

@@ -0,0 +1,30 @@
+from typing import Union
+
+from libbot.pyrogram.classes import PyroClient
+from pyrogram.types import User
+
+from classes.pyrouser import PyroUser
+from modules.database import col_users
+
+
+class PyroClient(PyroClient):
+    async def find_user(self, user: Union[int, User], group: int) -> PyroUser:
+        """Find User by it's ID or User object
+
+        ### Args:
+            * user (`Union[int, User]`): ID or User object to extract ID from
+            * group (`int`): ID of the group
+
+        ### Returns:
+            * `PyroUser`: PyroUser object
+        """
+        db_record = col_users.find_one(
+            {"id": user.id if isinstance(user, User) else user, "group": group}
+        )
+
+        if db_record is None:
+            raise KeyError(
+                f"User with ID {user.id if isinstance(user, User) else user} was not found in the database"
+            )
+
+        return PyroUser(**db_record)

classes/pyrouser.py

@@ -0,0 +1,104 @@
+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
+    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 = col_users.find_one({"id": id, "group": group})
+        if db_entry is None:
+            inserted = col_users.insert_one(
+                {
+                    "id": id,
+                    "group": group,
+                    "failed": failed,
+                    "emojis": emojis,
+                    "selected": selected,
+                    "score": score,
+                    "mistakes": mistakes,
+                }
+            )
+            db_entry = {
+                "_id": inserted.inserted_id,
+                "id": id,
+                "group": group,
+                "failed": failed,
+                "emojis": emojis,
+                "selected": selected,
+                "score": 0,
+                "mistakes": mistakes,
+            }
+        return cls(**db_entry)
+
+    def set_failed(self, failed: bool = True) -> None:
+        logger.debug("%s's failure state has been set to %s", self.id, failed)
+        col_users.update_one({"_id": self._id}, {"$set": {"failed": failed}})
+
+    def set_emojis(self, emojis: List[str]) -> None:
+        logger.debug("%s's emojis have been set to %s", self.id, emojis)
+        col_users.update_one({"_id": self._id}, {"$set": {"emojis": emojis}})
+
+    def set_score(self, score: int = 0) -> None:
+        logger.debug("%s's score has been set to %s", self.id, score)
+        col_users.update_one({"_id": self._id}, {"$set": {"score": score}})
+
+    def set_mistakes(self, mistakes: int = 0) -> None:
+        logger.debug("%s's mistakes count has been set to %s", self.id, mistakes)
+        col_users.update_one({"_id": self._id}, {"$set": {"mistakes": mistakes}})
+
+    def update_score(self, points: int = 1) -> None:
+        logger.debug("%s point(s) have been added to %s score", points, self.id)
+        col_users.update_one(
+            {"_id": self._id}, {"$set": {"score": self.score + points}}
+        )
+
+    def update_mistakes(self, points: int = 1) -> None:
+        logger.debug("%s point(s) have been added to %s mistakes", points, self.id)
+        col_users.update_one(
+            {"_id": self._id}, {"$set": {"mistakes": self.mistakes + points}}
+        )
+
+    def update_selected(self, entry: str) -> None:
+        logger.debug("Emoji %s has been added to %s's selection list", entry, self.id)
+        col_users.update_one(
+            {"_id": self._id}, {"$set": {"selected": self.selected + [entry]}}
+        )