From 032b636b025b56f1510b337cfd186b799962ac25 Mon Sep 17 00:00:00 2001 From: stijndcl Date: Thu, 30 Jun 2022 21:17:48 +0200 Subject: [PATCH] Nightly, bank, award & dinks --- ...> 0d03c226d881_initial_currency_models.py} | 8 +-- database/crud/currency.py | 43 +++++++++++++++ database/crud/users.py | 37 +++++++++++++ database/exceptions/currency.py | 2 + database/models.py | 3 +- didier/cogs/currency.py | 54 +++++++++++++++++++ didier/utils/discord/checks/__init__.py | 1 + .../utils/discord/checks/message_commands.py | 6 +++ didier/utils/types/string.py | 8 +++ 9 files changed, 157 insertions(+), 5 deletions(-) rename alembic/versions/{5f3a11a80e69_initial_currency_models.py => 0d03c226d881_initial_currency_models.py} (94%) create mode 100644 database/crud/currency.py create mode 100644 database/crud/users.py create mode 100644 database/exceptions/currency.py create mode 100644 didier/cogs/currency.py diff --git a/alembic/versions/5f3a11a80e69_initial_currency_models.py b/alembic/versions/0d03c226d881_initial_currency_models.py similarity index 94% rename from alembic/versions/5f3a11a80e69_initial_currency_models.py rename to alembic/versions/0d03c226d881_initial_currency_models.py index 96981b9..45a5e26 100644 --- a/alembic/versions/5f3a11a80e69_initial_currency_models.py +++ b/alembic/versions/0d03c226d881_initial_currency_models.py @@ -1,8 +1,8 @@ """Initial currency models -Revision ID: 5f3a11a80e69 +Revision ID: 0d03c226d881 Revises: b2d511552a1f -Create Date: 2022-06-30 19:40:02.701336 +Create Date: 2022-06-30 20:02:27.284759 """ from alembic import op @@ -10,7 +10,7 @@ import sqlalchemy as sa # revision identifiers, used by Alembic. -revision = '5f3a11a80e69' +revision = '0d03c226d881' down_revision = 'b2d511552a1f' branch_labels = None depends_on = None @@ -20,12 +20,12 @@ def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### op.create_table('users', sa.Column('user_id', sa.BigInteger(), nullable=False), - sa.Column('dinks', sa.BigInteger(), nullable=False), sa.PrimaryKeyConstraint('user_id') ) op.create_table('bank', sa.Column('bank_id', sa.Integer(), nullable=False), sa.Column('user_id', sa.BigInteger(), nullable=True), + sa.Column('dinks', sa.BigInteger(), nullable=False), sa.Column('interest_level', sa.Integer(), nullable=False), sa.Column('capacity_level', sa.Integer(), nullable=False), sa.Column('rob_level', sa.Integer(), nullable=False), diff --git a/database/crud/currency.py b/database/crud/currency.py new file mode 100644 index 0000000..9fe21e0 --- /dev/null +++ b/database/crud/currency.py @@ -0,0 +1,43 @@ +from datetime import datetime + +from sqlalchemy.ext.asyncio import AsyncSession + +from database.crud import users +from database.exceptions import currency as exceptions +from database.models import Bank + + +NIGHTLY_AMOUNT = 420 + + +async def get_bank(session: AsyncSession, user_id: int) -> Bank: + """Get a user's bank info""" + user = await users.get_or_add(session, user_id) + return user.bank + + +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""" + user = await users.get_or_add(session, user_id) + nightly_data = user.nightly_data + + now = datetime.now() + + if nightly_data.last_nightly is not None and nightly_data.last_nightly.date() == now.date(): + raise exceptions.DoubleNightly + + bank = user.bank + bank.dinks += NIGHTLY_AMOUNT + nightly_data.last_nightly = now + + session.add(bank) + session.add(nightly_data) + await session.commit() diff --git a/database/crud/users.py b/database/crud/users.py new file mode 100644 index 0000000..794d951 --- /dev/null +++ b/database/crud/users.py @@ -0,0 +1,37 @@ +from typing import Optional + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from database.models import User, Bank, NightlyData + + +async def get_or_add(session: AsyncSession, user_id: int) -> User: + """Get a user's profile + If it doesn't exist yet, create it (along with all linked datastructures). + """ + statement = select(User).where(User.user_id == user_id) + user: Optional[User] = (await session.execute(statement)).scalar_one_or_none() + + # User exists + if user is not None: + return user + + # Create new user + user = User(user_id=user_id) + session.add(user) + await session.commit() + + # Add bank & nightly info + bank = Bank(user_id=user_id) + nightly_data = NightlyData(user_id=user_id) + user.bank = bank + user.nightly_data = nightly_data + + session.add(bank) + session.add(nightly_data) + session.add(user) + + await session.commit() + + return user diff --git a/database/exceptions/currency.py b/database/exceptions/currency.py new file mode 100644 index 0000000..6eab7a1 --- /dev/null +++ b/database/exceptions/currency.py @@ -0,0 +1,2 @@ +class DoubleNightly(Exception): + """Exception raised when claiming nightlies multiple times per day""" diff --git a/database/models.py b/database/models.py index 02b16a5..0bc6370 100644 --- a/database/models.py +++ b/database/models.py @@ -17,6 +17,8 @@ class Bank(Base): bank_id: int = Column(Integer, primary_key=True) user_id: int = Column(BigInteger, ForeignKey("users.user_id")) + dinks: int = Column(BigInteger, default=0, nullable=False) + # Interest rate interest_level: int = Column(Integer, default=1, nullable=False) @@ -119,7 +121,6 @@ class User(Base): __tablename__ = "users" user_id: int = Column(BigInteger, primary_key=True) - dinks: int = Column(BigInteger, default=0, nullable=False) bank: Bank = relationship( "Bank", back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan" diff --git a/didier/cogs/currency.py b/didier/cogs/currency.py new file mode 100644 index 0000000..193a9ef --- /dev/null +++ b/didier/cogs/currency.py @@ -0,0 +1,54 @@ +import discord +from discord.ext import commands + +from didier import Didier + +from database.crud import currency as crud +from database.exceptions.currency import DoubleNightly +from didier.utils.discord.checks import is_owner +from didier.utils.types.string import pluralize + + +class Currency(commands.Cog): + """Everything Dinks-related""" + + client: Didier + + def __init__(self, client: Didier): + super().__init__() + self.client = client + + @commands.command(name="Award") + @commands.check(is_owner) + async def award(self, ctx: commands.Context, user: discord.User, amount: int): + async with self.client.db_session as session: + await crud.add_dinks(session, user.id, amount) + await self.client.confirm_message(ctx.message) + + @commands.hybrid_command(name="bank") + async def bank(self, ctx: commands.Context): + """Show your Didier Bank information""" + async with self.client.db_session as session: + await crud.get_bank(session, ctx.author.id) + + @commands.hybrid_command(name="dinks") + async def dinks(self, ctx: commands.Context): + """Check your Didier Dinks""" + async with self.client.db_session as session: + bank = await crud.get_bank(session, ctx.author.id) + plural = pluralize("Didier Dink", bank.dinks) + await ctx.reply(f"**{ctx.author.display_name}** heeft **{bank.dinks}** {plural}.", mention_author=False) + + @commands.hybrid_command(name="nightly") + async def nightly(self, ctx: commands.Context): + """Claim nightly Dinks""" + async with self.client.db_session as session: + try: + await crud.claim_nightly(session, ctx.author.id) + await ctx.reply(f"Je hebt je dagelijkse **{crud.NIGHTLY_AMOUNT}** Didier Dinks geclaimd.") + except DoubleNightly: + await ctx.reply("Je hebt je nightly al geclaimd vandaag.", mention_author=False, ephemeral=True) + + +async def setup(client: Didier): + await client.add_cog(Currency(client)) diff --git a/didier/utils/discord/checks/__init__.py b/didier/utils/discord/checks/__init__.py index e69de29..7332b42 100644 --- a/didier/utils/discord/checks/__init__.py +++ b/didier/utils/discord/checks/__init__.py @@ -0,0 +1 @@ +from .message_commands import is_owner diff --git a/didier/utils/discord/checks/message_commands.py b/didier/utils/discord/checks/message_commands.py index e69de29..a57b315 100644 --- a/didier/utils/discord/checks/message_commands.py +++ b/didier/utils/discord/checks/message_commands.py @@ -0,0 +1,6 @@ +from discord.ext import commands + + +async def is_owner(ctx: commands.Context) -> bool: + """Check that a command is being invoked by the owner of the bot""" + return await ctx.bot.is_owner(ctx.author) diff --git a/didier/utils/types/string.py b/didier/utils/types/string.py index eddfff8..0255877 100644 --- a/didier/utils/types/string.py +++ b/didier/utils/types/string.py @@ -20,3 +20,11 @@ def leading(character: str, string: str, target_length: Optional[int] = 2) -> st frequency = math.ceil((target_length - len(string)) / len(character)) return (frequency * character) + string + + +def pluralize(word: str, amount: int, plural_form: Optional[str] = None) -> str: + """Turn a word into plural""" + if amount == 1: + return word + + return plural_form or (word + "s")