didier/database/crud/currency.py

76 lines
1.9 KiB
Python

from datetime import date
from typing import Union
from sqlalchemy.ext.asyncio import AsyncSession
from database.crud import users
from database.exceptions import currency as exceptions
from database.schemas import Bank, NightlyData
__all__ = [
"add_dinks",
"claim_nightly",
"get_bank",
"get_nightly_data",
"invest",
"NIGHTLY_AMOUNT",
]
NIGHTLY_AMOUNT = 50
async def get_bank(session: AsyncSession, user_id: int) -> Bank:
"""Get a user's bank info"""
user = await users.get_or_add_user(session, user_id)
return user.bank
async def get_nightly_data(session: AsyncSession, user_id: int) -> NightlyData:
"""Get a user's nightly info"""
user = await users.get_or_add_user(session, user_id)
return user.nightly_data
async def invest(session: AsyncSession, user_id: int, amount: Union[str, int]) -> int:
"""Invest all your Dinks"""
bank = await get_bank(session, user_id)
if amount == "all":
amount = bank.dinks
# Don't allow investing more dinks than you own
amount = min(bank.dinks, int(amount))
bank.dinks -= amount
bank.invested += amount
session.add(bank)
await session.commit()
return amount
async def add_dinks(session: AsyncSession, user_id: int, amount: int):
"""Increase the Dinks counter for a user"""
bank = await get_bank(session, user_id)
bank.dinks += amount
session.add(bank)
await session.commit()
async def claim_nightly(session: AsyncSession, user_id: int):
"""Claim daily Dinks"""
nightly_data = await get_nightly_data(session, user_id)
now = date.today()
if nightly_data.last_nightly is not None and nightly_data.last_nightly == now:
raise exceptions.DoubleNightly
bank = await get_bank(session, user_id)
bank.dinks += NIGHTLY_AMOUNT
nightly_data.last_nightly = now
session.add(bank)
session.add(nightly_data)
await session.commit()