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 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 dosctring async def deposit(self, amount: float) -> None: await self._set("balance", round(self.balance + amount, 2)) # TODO Add a check to prevent negative balances # TODO Write a dosctring async def withdraw(self, amount: float) -> None: await self._set("balance", round(self.balance - amount, 2))