GloryBot/classes/pycord_member.py

import logging
from dataclasses import dataclass
from typing import Literal, Union

from bson import ObjectId
from discord import Thread

from classes.enums import MemberStatus
from classes.errors import MemberAlreadyExistsError, MemberNotFoundError
from classes.pycord_check import PycordCheck
from modules.database import col_members

logger = logging.getLogger(__name__)


@dataclass
class PycordMember:
    """Dataclass of DB entry of a member"""

    __slots__ = ("_id", "id", "guild", "status")

    _id: ObjectId
    id: int
    guild: ObjectId
    status: Literal[
        MemberStatus.VERIFIED,
        MemberStatus.UNVERIFIED,
        MemberStatus.FAILED,
        MemberStatus.ADDITIONAL,
    ]

    @classmethod
    async def find(cls, id: Union[int, ObjectId], guild: ObjectId):
        """Find member in the database.

        ### Args:
            * id (`Union[int, ObjectId]`): Member's Discord ID
            * guild (`ObjectId`): Discord guild's database ID

        ### Raises:
            * `MemberNotFoundError`: Raised when member entry could not be found.

        ### Returns:
            * `PycordMember`: Member with its database data.
        """
        db_entry = await col_members.find_one(
            (
                {"id": id, "guild": guild}
                if isinstance(id, int)
                else {"_id": id, "guild": guild}
            )
        )

        if not db_entry:
            raise MemberNotFoundError(id, guild)

        return cls(**db_entry)

    @classmethod
    async def create(
        cls,
        id: int,
        guild: ObjectId,
        status: Literal[
            MemberStatus.VERIFIED,
            MemberStatus.UNVERIFIED,
            MemberStatus.FAILED,
            MemberStatus.ADDITIONAL,
        ],
    ):
        # Check whether member already exists
        if await col_members.find_one({"id": True, "guild": guild}):
            raise MemberAlreadyExistsError(id, guild)

        # Create a member dict
        check = {
            "id": id,
            "guild": guild,
            "status": status.value,
        }

        # Insert the member into the database
        inserted = await col_members.insert_one(check)

        check["_id"] = inserted.inserted_id

        return cls(**check)

    async def new_check(self, thread_id: int) -> PycordCheck:
        return await PycordCheck.create(self.guild, thread_id, self._id)

    async def get_check(self) -> PycordCheck:
        """Get an ongoing check

        ### Returns:
            * `PycordCheck`: An ongoing check

        ### Raises:
            * `CheckNotFoundError`: Member does not have an ongoing check
        """
        return await PycordCheck.find(self.guild, self._id)

    async def set_status(
        self,
        status: Literal[
            MemberStatus.VERIFIED,
            MemberStatus.UNVERIFIED,
            MemberStatus.FAILED,
            MemberStatus.ADDITIONAL,
        ],
    ) -> Union[Thread, None]:
        await col_members.update_one({"_id": self._id}, {"$set": {status.value}})
        self.status = status