Bank upgrades

.gitignore

@@ -154,3 +154,6 @@ cython_debug/
 
 # PyCharm
 .idea/
+
+# Debugging files
+debug.py

database/crud/currency.py

@@ -5,6 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from database.crud import users
 from database.exceptions import currency as exceptions
 from database.models import Bank
+from database.utils.math.currency import rob_upgrade_price, interest_upgrade_price, capacity_upgrade_price
 
 
 NIGHTLY_AMOUNT = 420
@@ -41,3 +42,57 @@ async def claim_nightly(session: AsyncSession, user_id: int):
     session.add(bank)
     session.add(nightly_data)
     await session.commit()
+
+
+async def upgrade_capacity(database_session: AsyncSession, user_id: int) -> int:
+    """Upgrade capacity level"""
+    bank = await get_bank(database_session, user_id)
+    upgrade_price = capacity_upgrade_price(bank.capacity_level)
+
+    # Can't afford this upgrade
+    if upgrade_price > bank.dinks:
+        raise exceptions.NotEnoughDinks
+
+    bank.dinks -= upgrade_price
+    bank.capacity_level += 1
+
+    database_session.add(bank)
+    await database_session.commit()
+
+    return bank.capacity_level
+
+
+async def upgrade_interest(database_session: AsyncSession, user_id: int) -> int:
+    """Upgrade interest level"""
+    bank = await get_bank(database_session, user_id)
+    upgrade_price = interest_upgrade_price(bank.interest_level)
+
+    # Can't afford this upgrade
+    if upgrade_price > bank.dinks:
+        raise exceptions.NotEnoughDinks
+
+    bank.dinks -= upgrade_price
+    bank.interest_level += 1
+
+    database_session.add(bank)
+    await database_session.commit()
+
+    return bank.interest_level
+
+
+async def upgrade_rob(database_session: AsyncSession, user_id: int) -> int:
+    """Upgrade rob level"""
+    bank = await get_bank(database_session, user_id)
+    upgrade_price = rob_upgrade_price(bank.rob_level)
+
+    # Can't afford this upgrade
+    if upgrade_price > bank.dinks:
+        raise exceptions.NotEnoughDinks
+
+    bank.dinks -= upgrade_price
+    bank.rob_level += 1
+
+    database_session.add(bank)
+    await database_session.commit()
+
+    return bank.rob_level

database/exceptions/currency.py

@@ -1,2 +1,6 @@
 class DoubleNightly(Exception):
     """Exception raised when claiming nightlies multiple times per day"""
+
+
+class NotEnoughDinks(Exception):
+    """Exception raised when trying to do something you don't have the Dinks for"""

didier/cogs/currency.py

@@ -4,11 +4,11 @@ import discord
 from discord.ext import commands
 
 from database.crud import currency as crud
-from database.exceptions.currency import DoubleNightly
+from database.exceptions.currency import DoubleNightly, NotEnoughDinks
 from didier import Didier
 from didier.utils.discord.checks import is_owner
 from didier.utils.discord.converters import abbreviated_number
-from didier.utils.math.currency import capacity_upgrade_price, interest_upgrade_price, rob_upgrade_price
+from database.utils.math.currency import capacity_upgrade_price, interest_upgrade_price, rob_upgrade_price
 from didier.utils.types.string import pluralize
 
 
@@ -72,6 +72,39 @@ class Currency(commands.Cog):
 
         await ctx.reply(embed=embed, mention_author=False)
 
+    @bank_upgrades.command(name="Capacity", aliases=["C"])
+    async def bank_upgrade_capacity(self, ctx: commands.Context):
+        """Upgrade the capacity level of your bank"""
+        async with self.client.db_session as session:
+            try:
+                await crud.upgrade_capacity(session, ctx.author.id)
+                await ctx.message.add_reaction("⏫")
+            except NotEnoughDinks:
+                await ctx.reply("Je hebt niet genoeg Didier Dinks om dit te doen.", mention_author=False)
+                await self.client.reject_message(ctx.message)
+
+    @bank_upgrades.command(name="Interest", aliases=["I"])
+    async def bank_upgrade_interest(self, ctx: commands.Context):
+        """Upgrade the interest level of your bank"""
+        async with self.client.db_session as session:
+            try:
+                await crud.upgrade_interest(session, ctx.author.id)
+                await ctx.message.add_reaction("⏫")
+            except NotEnoughDinks:
+                await ctx.reply("Je hebt niet genoeg Didier Dinks om dit te doen.", mention_author=False)
+                await self.client.reject_message(ctx.message)
+
+    @bank_upgrades.command(name="Rob", aliases=["R"])
+    async def bank_upgrade_rob(self, ctx: commands.Context):
+        """Upgrade the rob level of your bank"""
+        async with self.client.db_session as session:
+            try:
+                await crud.upgrade_rob(session, ctx.author.id)
+                await ctx.message.add_reaction("⏫")
+            except NotEnoughDinks:
+                await ctx.reply("Je hebt niet genoeg Didier Dinks om dit te doen.", mention_author=False)
+                await self.client.reject_message(ctx.message)
+
     @commands.hybrid_command(name="dinks")
     async def dinks(self, ctx: commands.Context):
         """Check your Didier Dinks"""

pyproject.toml

@@ -30,7 +30,8 @@ disable = [
     "missing-module-docstring",
     "too-few-public-methods",
     "too-many-arguments",
-    "too-many-instance-attributes"
+    "too-many-instance-attributes",
+    "too-many-locals"
 ]
 
 [tool.pylint.format]