diff --git a/api/extensions/utils.py b/api/extensions/utils.py index efee25e..06e60f6 100644 --- a/api/extensions/utils.py +++ b/api/extensions/utils.py @@ -1,7 +1,9 @@ from pathlib import Path -from api.app import app + from fastapi.responses import FileResponse +from api.app import app + @app.get("/favicon.ico", response_class=FileResponse, include_in_schema=False) async def favicon(): diff --git a/classes/errors/__init__.py b/classes/errors/__init__.py index a293bb6..604e596 100644 --- a/classes/errors/__init__.py +++ b/classes/errors/__init__.py @@ -1,2 +1,7 @@ from .pycord_user import UserNotFoundError -from .wallet import WalletNotFoundError +from .wallet import ( + WalletBalanceLimitExceeded, + WalletInsufficientFunds, + WalletNotFoundError, + WalletOverdraftLimitExceeded, +) diff --git a/classes/errors/wallet.py b/classes/errors/wallet.py index a461cdb..28339ff 100644 --- a/classes/errors/wallet.py +++ b/classes/errors/wallet.py @@ -24,3 +24,39 @@ class WalletInsufficientFunds(Exception): 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})" ) + + +class WalletOverdraftLimitExceeded(Exception): + """Wallet's overdraft limit is not sufficient to perform the operation""" + + def __init__( + self, + wallet: "Wallet", + amount: float, + overdraft_limit: float, + ) -> None: + self.wallet = wallet + self.amount = amount + self.overdraft_limit = overdraft_limit + + 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}, overdraft limit: {self.overdraft_limit})" + ) + + +class WalletBalanceLimitExceeded(Exception): + """Wallet's balance limit is not high enough to perform the operation""" + + def __init__( + self, + wallet: "Wallet", + amount: float, + balance_limit: float, + ) -> None: + self.wallet = wallet + self.amount = amount + self.balance_limit = balance_limit + + super().__init__( + f"Wallet of a user with id {self.wallet.owner_id} for the guild with id {self.wallet.guild_id} would have too much funds after the operation (balance: {self.wallet.balance}, deposited: {self.amount}, balance limit: {self.balance_limit})" + ) diff --git a/classes/pycord_user.py b/classes/pycord_user.py index 835ef28..f6248ec 100644 --- a/classes/pycord_user.py +++ b/classes/pycord_user.py @@ -82,9 +82,7 @@ class PycordUser: setattr(self, key, value) - await col_users.update_one( - {"_id": self._id}, {"$set": {key: value}}, upsert=True - ) + await col_users.update_one({"_id": self._id}, {"$set": {key: value}}, upsert=True) self._update_cache(cache) diff --git a/classes/wallet.py b/classes/wallet.py index 5e38788..6302335 100644 --- a/classes/wallet.py +++ b/classes/wallet.py @@ -7,7 +7,11 @@ from typing import Any, Dict, Optional from bson import ObjectId from pymongo.results import InsertOneResult -from classes.errors import WalletNotFoundError +from classes.errors import ( + WalletBalanceLimitExceeded, + WalletNotFoundError, + WalletOverdraftLimitExceeded, +) from classes.errors.wallet import WalletInsufficientFunds from modules.database import col_wallets @@ -28,9 +32,7 @@ class Wallet: 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} - ) + db_entry = await col_wallets.find_one({"owner_id": owner_id, "guild_id": guild_id}) if db_entry is None: if not allow_creation: @@ -60,13 +62,9 @@ class Wallet: setattr(self, key, value) - await col_wallets.update_one( - {"_id": self._id}, {"$set": {key: value}}, upsert=True - ) + 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 - ) + logger.info("Set attribute '%s' of the wallet %s to '%s'", key, str(self._id), value) @staticmethod def get_defaults( @@ -96,8 +94,15 @@ class Wallet: 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)) + 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( @@ -105,10 +110,47 @@ class Wallet: 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) + ) -> float: + if amount > self.balance: + if not allow_overdraft or overdraft_limit is None: + raise WalletInsufficientFunds(self, amount) - await self._set("balance", round(self.balance - amount, 2)) + 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) diff --git a/migrate.py b/migrate.py index 25a99f6..85684b0 100644 --- a/migrate.py +++ b/migrate.py @@ -1,4 +1,3 @@ from modules.migrator import migrate_database - migrate_database() diff --git a/modules/database.py b/modules/database.py index 08f161f..8fb6557 100644 --- a/modules/database.py +++ b/modules/database.py @@ -35,6 +35,4 @@ col_wallets: AsyncCollection = db.get_collection("wallets") # Update indexes db.dispatch.get_collection("users").create_index("id", unique=True) -db.dispatch.get_collection("wallets").create_index( - ["owner_id", "guild_id"], unique=False -) +db.dispatch.get_collection("wallets").create_index(["owner_id", "guild_id"], unique=False) diff --git a/modules/extensions_loader.py b/modules/extensions_loader.py index 64cabd6..c50b0e6 100644 --- a/modules/extensions_loader.py +++ b/modules/extensions_loader.py @@ -15,9 +15,7 @@ def get_py_files(src: str | Path) -> List[str]: py_files = [] for root, dirs, files in walk(src): - py_files.extend( - Path(f"{cwd}/{root}/{file}") for file in files if file.endswith(".py") - ) + py_files.extend(Path(f"{cwd}/{root}/{file}") for file in files if file.endswith(".py")) return py_files diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4fe7567 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,6 @@ +[tool.black] +line-length = 96 +target-version = ["py311"] + +[tool.isort] +profile = "black" \ No newline at end of file