Merge pull request #233 from sarzz2/master

added persistent role,levelling features

bot/extensions/levelling/__init__.py

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

bot/extensions/levelling/commands.py

@@ -0,0 +1,156 @@
+import datetime
+import random
+
+import asyncpg.exceptions
+import discord
+from discord import app_commands
+from discord.ext import commands
+
+from bot import core
+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]
+        self.xp_boost = 1
+
+    async def cog_load(self):
+        for guild in self.bot.guilds:
+            data = await IgnoredChannel.list_by_guild(guild_id=guild.id)
+            for i in data:
+                if guild.id not in self.ignored_channel:
+                    self.ignored_channel[guild.id] = [i.channel_id]
+                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
+        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) * self.xp_boost
+        # Add the XP and update the DB
+        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 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("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}",
+                    timestamp=datetime.datetime.utcnow(),
+                    colour=discord.Colour.blurple(),
+                )
+                embed.set_thumbnail(url=member.avatar)
+                return await interaction.response.send_message(embed=embed)
+
+    @app_commands.command()
+    @app_commands.checks.has_permissions(administrator=True)
+    async def ignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel):
+        """Add the channel to the ignored channel list to not gain XP"""
+        await IgnoredChannel.insert_by_guild(channel.guild.id, channel.id)
+        await interaction.response.send_message(f"{channel} has been ignored from gaining XP.")
+
+    @app_commands.command()
+    @app_commands.checks.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)
+        await interaction.response.send_message(f"{channel} has been removed from ignored channel list")
+
+    @app_commands.command()
+    @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"""
+        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"""
+        if xp <= 0:
+            return await interaction.response.send_message("XP can not be less than 0")
+        try:
+            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")
+
+    @app_commands.command()
+    async def levelling_rewards_list(self, interaction: core.InteractionType):
+        data = await LevellingRole.list_by_guild(interaction.guild.id)
+        res = "| {:<10} | {:<10} | {:<5} |".format("Guild ID", "Role ID", "Level")
+        res += "\n|" + "-" * 12 + "|" + "-" * 12 + "|" + "-" * 7 + "|"
+
+        # Print data
+        for record in data:
+            res += "\n| {:<10} | {:<10} | {:<5} |".format(record["guild_id"], record["role_id"], record["level"])
+
+        await interaction.response.send_message(f"{res}")
+
+    @app_commands.command()
+    @app_commands.checks.has_permissions(administrator=True)
+    async def xp_multiplier(self, interaction: core.InteractionType, multiplier: int):
+        """Increase XP gain per message"""
+        if multiplier <= 0 or multiplier > 5:
+            return await interaction.response.send_message("Invalid multiplier value.(Max. 5)")
+
+        self.xp_boost = multiplier
+        return await interaction.response.send_message(f"XP multiplied by {multiplier}x.")
+
+
+async def setup(bot: commands.Bot):
+    await bot.add_cog(Levelling(bot=bot))

bot/extensions/levelling/events.py

@@ -0,0 +1,43 @@
+from bisect import bisect
+
+import discord
+from discord.ext import commands
+
+from bot import core
+from bot.models import LevellingRole, PersistentRole
+
+
+class LevelEvents(commands.Cog):
+    """Events for Levelling in discord."""
+
+    def __init__(self, bot: core.DiscordBot):
+        self.bot = bot
+
+    @commands.Cog.listener()
+    async def on_level_up(self, new_level, member: discord.Member):
+        """Add roles when user levels up"""
+        data = await LevellingRole.list_by_guild(guild_id=member.guild.id)
+        for i in range(len(data)):
+            # Adding roles when user levels up
+            if new_level >= data[i].level and member.guild.get_role(data[i].role_id) not in member.roles:
+                await member.add_roles(member.guild.get_role(data[i].role_id))
+                await PersistentRole.insert_by_guild(
+                    guild_id=member.guild.id, user_id=member.id, role_id=data[i].role_id
+                )
+            # Removing roles when users level down using remove_xp command
+            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/extensions/persistent_roles/events.py

@@ -0,0 +1,24 @@
+from discord.ext import commands
+
+from bot import core
+from bot.models import PersistentRole
+
+
+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
+        if data is None:
+            return
+        # Add the roles to the user if data exists
+        else:
+            for i in range(len(data)):
+                await member.add_roles(member.guild.get_role(data[i].role_id))

bot/models/__init__.py

@@ -1,15 +1,23 @@
 from .gconfig import FilterConfig
+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 .persisted_role import PersistentRole
 from .rep import Rep
 from .tag import Tag
 from .user import User
 
-__all__ = (  # Fixes F401
+__all__ = (
     Model,
     FilterConfig,
     Message,
     Rep,
     Tag,
     User,
-)
+    Levels,
+    PersistentRole,
+    IgnoredChannel,
+    LevellingRole,
+)  # Fixes F401

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/levelling_ignored_channels.py

@@ -0,0 +1,24 @@
+from .model import Model
+
+
+class IgnoredChannel(Model):
+    id: int
+    guild_id: int
+    channel_id: int
+
+    @classmethod
+    async def list_by_guild(cls, guild_id: int):
+        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 levelling_ignored_channels (guild_id,channel_id) VALUES($1, $2)
+                   ON CONFLICT (guild_id,channel_id) DO NOTHING
+                     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 levelling_ignored_channels WHERE guild_id= $1 and channel_id = $2"""
+        return await cls.fetch(query, guild_id, channel_id)

bot/models/levelling_roles.py

@@ -0,0 +1,26 @@
+from .model import Model
+
+
+class LevellingRole(Model):
+    id: int
+    guild_id: int
+    role_id: int
+    level: int
+
+    @classmethod
+    async def list_by_guild(cls, guild_id: int):
+        query = """SELECT * FROM levelling_roles WHERE guild_id=$1"""
+        return await cls.fetch(query, guild_id)
+
+    @classmethod
+    async def insert_by_guild(cls, guild_id: int, role_id: int, level: int):
+        query = """INSERT INTO levelling_roles (guild_id, role_id, level)
+                        VALUES ($1, $2, $3)
+                   ON CONFLICT (guild_id, role_id)
+                    DO NOTHING"""
+        await cls.execute(query, guild_id, role_id, level)
+
+    @classmethod
+    async def delete_by_guild(cls, guild_id: int, role_id: int):
+        query = """DELETE FROM levelling_roles WHERE guild_id = $1 and role_id = $2"""
+        await cls.execute(query, guild_id, role_id)