115 lines
3.2 KiB
Python
115 lines
3.2 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 import WalletNotFoundError
|
|
from classes.errors.wallet import 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) -> Dict[str, Any]:
|
|
return {
|
|
"_id": 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) -> None:
|
|
await self._set("balance", round(self.balance + amount, 2))
|
|
|
|
# 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))
|