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.wallet import ( WalletBalanceLimitExceeded, WalletNotFoundError, WalletOverdraftLimitExceeded, 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, json_compatible: bool = False) -> Dict[str, Any]: """Convert Wallet object to a JSON representation. Args: json_compatible (bool): Whether the JSON-incompatible objects like ObjectId need to be converted Returns: Dict[str, Any]: JSON representation of Wallet """ return { "_id": self._id if not json_compatible else str(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, balance_limit: Optional[float] = None) -> float: new_balance: float = round(self.balance + amount, 2) if balance_limit is not None and new_balance > balance_limit: raise WalletBalanceLimitExceeded(self, amount, balance_limit) await self._set("balance", new_balance) return new_balance # TODO Write a docstring async def withdraw( self, amount: float, allow_overdraft: bool = False, overdraft_limit: Optional[float] = None, ) -> float: if amount > self.balance: if not allow_overdraft or overdraft_limit is None: raise WalletInsufficientFunds(self, amount) if allow_overdraft and amount > overdraft_limit: raise WalletOverdraftLimitExceeded(self, amount, overdraft_limit) new_balance: float = round(self.balance - amount, 2) await self._set("balance", new_balance) return new_balance async def transfer( self, wallet_owner_id: int, wallet_guild_id: int, amount: float, balance_limit: Optional[float] = None, allow_overdraft: bool = False, overdraft_limit: Optional[float] = None, ) -> None: # TODO Replace with a concrete exception if amount < 0: raise ValueError() wallet: Wallet = await self.from_id( wallet_owner_id, wallet_guild_id, allow_creation=False ) if balance_limit is not None and amount + wallet.balance > balance_limit: raise WalletBalanceLimitExceeded(wallet, amount, balance_limit) if amount > self.balance: if not allow_overdraft or overdraft_limit is None: raise WalletInsufficientFunds(self, amount) if allow_overdraft and amount > overdraft_limit: raise WalletOverdraftLimitExceeded(self, amount, overdraft_limit) # TODO Make a sanity check to revert the transaction if anything goes wrong await self.withdraw(amount, allow_overdraft, overdraft_limit) await wallet.deposit(amount, balance_limit)