Javelina/classes/wallet.py

import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from logging import Logger
from typing import Any, Dict, Optional

from bson import ObjectId
from pymongo.results import InsertOneResult

from classes.errors import WalletNotFoundError
from classes.errors.wallet import WalletInsufficientFunds
from modules.database import col_wallets

logger: Logger = logging.getLogger(__name__)


@dataclass
class Wallet:
    _id: ObjectId
    owner_id: int
    guild_id: int
    balance: float
    is_frozen: bool
    created: datetime

    # TODO Write a docstring
    @classmethod
    async def from_id(
            cls, owner_id: int, guild_id: int, allow_creation: bool = True
    ) -> "Wallet":
        db_entry = await col_wallets.find_one(
            {"owner_id": owner_id, "guild_id": guild_id}
        )

        if db_entry is None:
            if not allow_creation:
                raise WalletNotFoundError(owner_id, guild_id)

            db_entry = Wallet.get_defaults(owner_id, guild_id)

            insert_result: InsertOneResult = await col_wallets.insert_one(db_entry)

            db_entry["_id"] = insert_result.inserted_id

        return cls(**db_entry)

    def _to_dict(self) -> Dict[str, Any]:
        return {
            "_id": self._id,
            "owner_id": self.owner_id,
            "guild_id": self.guild_id,
            "balance": self.balance,
            "is_frozen": self.is_frozen,
            "created": self.created,
        }

    async def _set(self, key: str, value: Any) -> None:
        if not hasattr(self, key):
            raise AttributeError()

        setattr(self, key, value)

        await col_wallets.update_one(
            {"_id": self._id}, {"$set": {key: value}}, upsert=True
        )

        logger.info(
            "Set attribute '%s' of the wallet %s to '%s'", key, str(self._id), value
        )

    @staticmethod
    def get_defaults(
            owner_id: Optional[int] = None, guild_id: Optional[int] = None
    ) -> Dict[str, Any]:
        return {
            "owner_id": owner_id,
            "guild_id": guild_id,
            "balance": 0.0,
            "is_frozen": False,
            "created": datetime.now(tz=timezone.utc),
        }

    @staticmethod
    def get_default_value(key: str) -> Any:
        if key not in Wallet.get_defaults():
            raise KeyError(f"There's no default value for key '{key}' in Wallet")

        return Wallet.get_defaults()[key]

    # TODO Write a docstring
    async def freeze(self) -> None:
        await self._set("is_frozen", True)

    # TODO Write a docstring
    async def unfreeze(self) -> None:
        await self._set("is_frozen", False)

    # TODO Write a docstring
    async def deposit(self, amount: float) -> None:
        await self._set("balance", round(self.balance + amount, 2))

    # TODO Write a docstring
    async def withdraw(
            self,
            amount: float,
            allow_overdraft: bool = False,
            overdraft_limit: Optional[float] = None,
    ) -> None:
        if amount > self.balance and (
                not allow_overdraft or (allow_overdraft and amount > overdraft_limit)
        ):
            raise WalletInsufficientFunds(self, amount)

        await self._set("balance", round(self.balance - amount, 2))