165 lines
5.1 KiB
Python
165 lines
5.1 KiB
Python
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)
|