diff --git a/classes/errors/wallet.py b/classes/errors/wallet.py index 3e94916..5593c3b 100644 --- a/classes/errors/wallet.py +++ b/classes/errors/wallet.py @@ -8,3 +8,19 @@ class WalletNotFoundError(Exception): super().__init__( f"Wallet of a user with id {self.owner_id} was not found for the guild with id {self.guild_id}" ) + + +class WalletInsufficientFunds(Exception): + """Wallet's balance is not sufficient to perform the operation""" + + def __init__( + self, + wallet: "Wallet", + amount: float, + ) -> None: + self.wallet = wallet + self.amount = amount + + super().__init__( + f"Wallet of a user with id {self.wallet.owner_id} for the guild with id {self.wallet.guild_id} does not have sufficient funds to perform the operation (balance: {self.wallet.balance}, requested: {self.amount})" + ) diff --git a/classes/wallet.py b/classes/wallet.py index 19d6bb7..ade5f7e 100644 --- a/classes/wallet.py +++ b/classes/wallet.py @@ -8,6 +8,7 @@ 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__) @@ -25,7 +26,7 @@ class Wallet: # TODO Write a docstring @classmethod async def from_id( - cls, owner_id: int, guild_id: int, allow_creation: bool = True + 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} @@ -69,7 +70,7 @@ class Wallet: @staticmethod def get_defaults( - owner_id: Optional[int] = None, guild_id: Optional[int] = None + owner_id: Optional[int] = None, guild_id: Optional[int] = None ) -> Dict[str, Any]: return { "owner_id": owner_id, @@ -94,11 +95,20 @@ class Wallet: async def unfreeze(self) -> None: await self._set("is_frozen", False) - # TODO Write a dosctring + # TODO Write a docstring 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: + # 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))