Added withdrawals and deposits

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})"
+        )

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))