refactored code, fixed models and added new commands for levelling role

bot/extensions/levelling/commands.py

@@ -1,123 +1,21 @@
 import datetime
 import random
 
+import asyncpg.exceptions
 import discord
 from discord import app_commands
 from discord.ext import commands
-from discord.ext.commands import has_permissions
 
 from bot import core
-from bot.models import IgnoredChannel, Levels
-
-REQUIRED_XP = [
-    0,
-    100,
-    255,
-    475,
-    770,
-    1150,
-    1625,
-    2205,
-    2900,
-    3720,
-    4675,
-    5775,
-    7030,
-    8450,
-    10045,
-    11825,
-    13800,
-    15980,
-    18375,
-    20995,
-    23850,
-    26950,
-    30305,
-    33925,
-    37820,
-    42000,
-    46475,
-    51255,
-    56350,
-    61770,
-    67525,
-    73625,
-    80080,
-    86900,
-    94095,
-    101675,
-    109650,
-    118030,
-    126825,
-    136045,
-    145700,
-    155800,
-    166355,
-    177375,
-    188870,
-    200850,
-    213325,
-    226305,
-    239800,
-    253820,
-    268375,
-    283475,
-    299130,
-    315350,
-    332145,
-    349525,
-    367500,
-    386080,
-    405275,
-    425095,
-    445550,
-    466650,
-    488405,
-    510825,
-    533920,
-    557700,
-    582175,
-    607355,
-    633250,
-    659870,
-    687225,
-    715325,
-    744180,
-    773800,
-    804195,
-    835375,
-    867350,
-    900130,
-    933725,
-    968145,
-    1003400,
-    1039500,
-    1076455,
-    1114275,
-    1152970,
-    1192550,
-    1233025,
-    1274405,
-    1316700,
-    1359920,
-    1404075,
-    1449175,
-    1495230,
-    1542250,
-    1590245,
-    1639225,
-    1689200,
-    1740180,
-    1899250,
-    1792175,
-    1845195,
-]
+from bot.models import IgnoredChannel, LevellingRole, Levels
+from bot.models.custom_roles import CustomRoles
 
 
 class Levelling(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
         self.ignored_channel = {}
+        self.required_xp = [0]
 
     async def cog_load(self):
         for guild in self.bot.guilds:
@@ -128,40 +26,46 @@ async def cog_load(self):
                 else:
                     self.ignored_channel[guild.id].append(i.channel_id)
 
+        # Calculating required_XP for next level and storing in a list, list-index corresponds to the level
+        for lvl in range(101):
+            xp = 5 * (lvl**2) + (50 * lvl) + 100
+            self.required_xp.append(xp + self.required_xp[-1])
+
     @commands.Cog.listener()
     async def on_message(self, message):
         # Return if message was sent by Bot or sent in DMs
         if message.author.bot or message.guild is None:
             return
 
         # Check if message is sent in ignored channel
-        if message.channel.id in self.ignored_channel[message.guild.id]:
-            return
+        try:
+            if message.channel.id in self.ignored_channel[message.guild.id]:
+                return
+        except KeyError:
+            pass
 
         # Generate random XP to be added
         xp = random.randint(5, 25)
         # Add the XP and update the DB
-        data = await Levels.update(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp)
-        await self._check_level_up(data, message.author)
+        data = await Levels.insert_by_guild(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp)
+        self.bot.dispatch("xp_updated", data=data, member=message.author, required_xp=self.required_xp)
 
     @app_commands.command()
     async def rank(self, interaction: core.InteractionType, member: discord.Member = None):
         """Check the rank of another member or yourself"""
         if member is None:
             member = interaction.user
-
         query = """WITH ordered_users AS (
-                        SELECT
-                        user_id,guild_id,total_xp,
-                        ROW_NUMBER() OVER (ORDER BY levels.total_xp DESC) AS rank
-                        FROM levels
-                        WHERE guild_id = $2)
-                        SELECT rank,total_xp,user_id,guild_id
-                        FROM ordered_users WHERE ordered_users.user_id = $1;"""
+                 SELECT id,user_id,guild_id,total_xp,
+            ROW_NUMBER() OVER (ORDER BY levelling_users.total_xp DESC) AS rank
+                    FROM levelling_users
+                    WHERE guild_id = $2)
+                   SELECT id,rank,total_xp,user_id,guild_id
+                     FROM ordered_users WHERE ordered_users.user_id = $1;"""
         data = await Levels.fetchrow(query, member.id, member.guild.id)
         if data is None:
-            return await interaction.response.send_message("You are not ranked yet!", ephemeral=True)
-        for level, j in enumerate(REQUIRED_XP):
+            return await interaction.response.send_message("User Not ranked yet!", ephemeral=True)
+        for level, j in enumerate(self.required_xp):
             if data.total_xp <= j:
                 embed = discord.Embed(
                     title=f"Rank: {data.rank}\nLevel: {level - 1}\nTotal XP:{data.total_xp}",
@@ -180,7 +84,6 @@ async def ignore_channel(self, interaction: core.InteractionType, channel: disco
 
     @app_commands.command()
     @app_commands.checks.has_permissions(administrator=True)
-    @has_permissions(administrator=True)
     async def unignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel):
         """Remove channel from ignored channel list"""
         await IgnoredChannel.delete_by_guild(channel.guild.id, channel.id)
@@ -190,35 +93,40 @@ async def unignore_channel(self, interaction: core.InteractionType, channel: dis
     @app_commands.checks.has_permissions(administrator=True)
     async def give_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member):
         """Give XP to specific user"""
-        data = await Levels.update(member.guild.id, member.id, xp)
-        await self._check_level_up(data, member)
-        await interaction.response.send_message(f"{xp} XP has been added to user {member}")
+        if xp <= 0:
+            return await interaction.response.send_message("XP can not be less than 0")
+        try:
+            data = await Levels.give_xp(member.guild.id, member.id, xp)
+            self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp)
+            await interaction.response.send_message(f"{xp} XP has been added to user {member}")
+        except asyncpg.exceptions.DataError:
+            return await interaction.response.send_message("Invalid XP provided")
 
     @app_commands.command()
     @app_commands.checks.has_permissions(administrator=True)
     async def remove_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member):
         """Remove XP from user"""
-        data = await Levels.remove_xp(member.guild.id, member.id, xp)
-        await self._check_level_up(data, member)
-        await interaction.response.send_message(f"{xp} XP has been removed from user {member}")
-
-    async def _check_level_up(self, data, member: discord.Member):
-        """Function to check if user's level has changed and trigger the event to assign the roles"""
-        # Calculating old and new level
+        if xp <= 0:
+            return await interaction.response.send_message("XP can not be less than 0")
         try:
-            for level, j in enumerate(REQUIRED_XP):
-                if data.old_total_xp <= j:
-                    old_level = level - 1
-                    break
-            for level, j in enumerate(REQUIRED_XP):
-                if data.total_xp <= j:
-                    new_level = level - 1
-                    break
-            # If the level has changed call level_up to handle roles change
-            if old_level != new_level:
-                self.bot.dispatch("level_up", new_level=new_level, member=member)
-        except AttributeError:
-            pass
+            data = await Levels.remove_xp(member.guild.id, member.id, xp)
+            self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp)
+            await interaction.response.send_message(f"{xp} XP has been removed from user {member}")
+        except asyncpg.exceptions.DataError:
+            return await interaction.response.send_message("Invalid XP provided")
+
+    @app_commands.command()
+    @app_commands.checks.has_permissions(administrator=True)
+    async def levelling_rewards_add(self, interaction: core.InteractionType, role: discord.Role, level: int):
+        await CustomRoles.insert_by_guild(role.id, role.guild.id, role.name, str(role.colour))
+        await LevellingRole.insert_by_guild(role.guild.id, role.id, level)
+        return await interaction.response.send_message("Levelling reward role added")
+
+    @app_commands.command()
+    @app_commands.checks.has_permissions(administrator=True)
+    async def levelling_rewards_remove(self, interaction: core.InteractionType, role: discord.Role):
+        await LevellingRole.delete_by_guild(role.guild.id, role.id)
+        return await interaction.response.send_message("Levelling reward role removed")
 
 
 async def setup(bot: commands.Bot):

bot/extensions/levelling/events.py

@@ -1,3 +1,5 @@
+from bisect import bisect
+
 import discord
 from discord.ext import commands
 
@@ -26,3 +28,16 @@ async def on_level_up(self, new_level, member: discord.Member):
             elif new_level <= data[i].level and member.guild.get_role(data[i].role_id) in member.roles:
                 await member.remove_roles(member.guild.get_role(data[i].role_id))
                 await PersistentRole.delete_by_guild(guild_id=member.guild.id, user_id=member.id)
+
+    @commands.Cog.listener()
+    async def on_xp_updated(self, data, member: discord.Member, required_xp):
+        """Function to check if user's level has changed and trigger the event to assign the roles"""
+        # Calculating old and new level
+        try:
+            old_level = bisect(required_xp, data.old_total_xp) - 1
+            new_level = bisect(required_xp, data.total_xp) - 1
+
+            if old_level != new_level:
+                self.bot.dispatch("level_up", new_level=new_level, member=member)
+        except AttributeError:
+            pass

bot/extensions/persistent_roles/__init__.py

@@ -0,0 +1,9 @@
+from bot.core import DiscordBot
+
+from .commands import PersistentRoles
+from .events import PersistentEvents
+
+
+async def setup(bot: DiscordBot) -> None:
+    await bot.add_cog(PersistentRoles(bot=bot))
+    await bot.add_cog(PersistentEvents(bot=bot))

bot/extensions/persistent_roles/commands.py

@@ -0,0 +1,10 @@
+from discord.ext import commands
+
+
+class PersistentRoles(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+
+async def setup(bot: commands.Bot):
+    await bot.add_cog(PersistentRoles(bot=bot))

bot/cogs/persistent_roles.py → bot/extensions/persistent_roles/events.py

@@ -1,14 +1,18 @@
 from discord.ext import commands
 
+from bot import core
 from bot.models import PersistentRole
 
 
-class PersistentRoles(commands.Cog):
-    def __init__(self, bot):
+class PersistentEvents(commands.Cog):
+    """Events for Persisted roles."""
+
+    def __init__(self, bot: core.DiscordBot):
         self.bot = bot
 
     @commands.Cog.listener()
     async def on_member_join(self, member):
+        """Add the persisted role to users if any on member join"""
         # Get the data
         data = await PersistentRole.list_by_guild(member.guild.id, member.id)
         # Return if data for specified user and guild does not exist
@@ -18,7 +22,3 @@ async def on_member_join(self, member):
         else:
             for i in range(len(data)):
                 await member.add_roles(member.guild.get_role(data[i].role_id))
-
-
-async def setup(bot: commands.Bot):
-    await bot.add_cog(PersistentRoles(bot=bot))

bot/models/__init__.py

@@ -1,10 +1,10 @@
 from .gconfig import FilterConfig
-from .ignored_channel import IgnoredChannel
-from .levelling_role import LevellingRole
-from .levels import Levels
+from .levelling_ignored_channels import IgnoredChannel
+from .levelling_roles import LevellingRole
+from .levelling_users import Levels
 from .message import Message
 from .model import Model
-from .persistent_role import PersistentRole
+from .persisted_role import PersistentRole
 from .rep import Rep
 from .tag import Tag
 from .user import User

bot/models/custom_roles.py

@@ -0,0 +1,20 @@
+from .model import Model
+
+
+class CustomRoles(Model):
+    id: int
+    role_id: int
+    guild_id: int
+    name: str
+    color: str
+
+    @classmethod
+    async def insert_by_guild(cls, role_id: int, guild_id: int, name: str, color: str):
+        query = """INSERT INTO custom_roles (role_id, guild_id, name, color)
+                        VALUES($1, $2, $3, $4)"""
+        await cls.execute(query, role_id, guild_id, name, color)
+
+    @classmethod
+    async def delete_by_guild(cls, guild_id: int, role_id: int):
+        query = """DELETE FROM custom_roles WHERE guild_id = $1 and role_id = $2"""
+        await cls.execute(query, guild_id, role_id)

bot/models/ignored_channel.py → bot/models/levelling_ignored_channels.py

@@ -2,22 +2,23 @@
 
 
 class IgnoredChannel(Model):
+    id: int
     guild_id: int
     channel_id: int
 
     @classmethod
     async def list_by_guild(cls, guild_id: int):
-        query = """SELECT * FROM ignored_channels WHERE guild_id = $1"""
+        query = """SELECT * FROM levelling_ignored_channels WHERE guild_id = $1"""
         return await cls.fetch(query, guild_id)
 
     @classmethod
     async def insert_by_guild(cls, guild_id: int, channel_id: int):
-        query = """INSERT INTO ignored_channels (guild_id,channel_id) VALUES($1, $2)
+        query = """INSERT INTO levelling_ignored_channels (guild_id,channel_id) VALUES($1, $2)
                    ON CONFLICT (guild_id,channel_id) DO NOTHING
-                     RETURNING guild_id, channel_id"""
+                     RETURNING id, guild_id, channel_id"""
         return await cls.fetch(query, guild_id, channel_id)
 
     @classmethod
     async def delete_by_guild(cls, guild_id: int, channel_id: int):
-        query = """DELETE FROM ignored_channel WHERE guild_id= $1 and channel_id = $2"""
+        query = """DELETE FROM levelling_ignored_channels WHERE guild_id= $1 and channel_id = $2"""
         return await cls.fetch(query, guild_id, channel_id)