Migrated poll cog to new extension format

bot/extensions/levelling/commands.py

@@ -230,14 +230,13 @@ def convert_int(integer):
         xp_offset_x -= xp_text_size[2] - xp_text_size[0]
         draw.text((xp_offset_x, xp_offset_y), text, font=self.small_font, fill="#fff")
 
-        if len(username) >= 15:
+        if len(username) >= 18:
             # Truncating the name
             username = username[:15] + "..."
 
         text_bbox = draw.textbbox((0, 0), username, font=self.medium_font)
-        text_offset_x = bar_offset_x - 10
         text_offset_y = bar_offset_y - (text_bbox[3] - text_bbox[1]) - 20
-        draw.text((text_offset_x, text_offset_y), username, font=self.medium_font, fill="#fff")
+        draw.text((bar_offset_x, text_offset_y), username, font=self.medium_font, fill="#fff")
 
         # create copy of background
         background = self.background.copy()
@@ -274,8 +273,6 @@ async def rank(self, interaction: core.InteractionType, member: discord.Member =
 
         record = await LevellingUser.pool.fetchrow(query, interaction.guild.id, member.id)
 
-        log.info(record)
-
         if record.total_xp is None:
             if member.id == interaction.user.id:
                 return await interaction.response.send_message(

bot/extensions/polls/__init__.py

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

bot/cogs/poll.py → bot/extensions/polls/commands.py

@@ -1,8 +1,12 @@
 import discord
+from discord import app_commands
 from discord.ext import commands
 
+from bot import core
+from bot.extensions.polls.views import PollButtons
 
-class Polls(commands.Cog):
+
+class Polls(commands.GroupCog, group_name="poll"):
     def __init__(self, bot: commands.AutoShardedBot):
         self.__bot = bot
 
@@ -30,80 +34,39 @@ def poll_check(self, message: discord.Message):
             return message.author == self.__bot.user
         return False
 
-    @commands.Cog.listener()
-    async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent):
-        channel: discord.TextChannel = self.__bot.get_channel(payload.channel_id)
-        message: discord.Message = await channel.fetch_message(payload.message_id)
-
-        if payload.user_id == self.__bot.user.id:
-            return
-
-        if not self.poll_check(message):
-            return
-
-        emojis = list(self.reactions.values())
-        if str(payload.emoji) not in emojis:
-            return
-
-        for reaction in message.reactions:
-            if str(reaction) not in emojis:
-                return
-
-            if str(reaction.emoji) != str(payload.emoji):
-                user = self.__bot.get_user(payload.user_id)
-                await message.remove_reaction(reaction.emoji, user)
-
-    @commands.group()
-    async def poll(self, ctx):
-        """Polls"""
-        if ctx.invoked_subcommand is None:
-            return await ctx.send_help(self.__bot.get_command("poll"))
-
-    @poll.command()
-    @commands.cooldown(1, 10, commands.BucketType.channel)
-    async def new(self, ctx, desc: str, *choices):
+    @app_commands.command()
+    @app_commands.checks.cooldown(1, 10)
+    async def new(self, interaction: core.InteractionType, title: str):
         """Create a new poll"""
-        await ctx.message.delete()
-
-        if len(choices) < 2:
-            ctx.command.reset_cooldown(ctx)
-            if len(choices) == 1:
-                return await ctx.send("Can't make a poll with only one choice")
-            return await ctx.send("You have to enter two or more choices to make a poll")
-
-        if len(choices) > 10:
-            ctx.command.reset_cooldown(ctx)
-            return await ctx.send("You can't make a poll with more than 10 choices")
 
         embed = discord.Embed(
-            description=f"**{desc}**\n\n"
-            + "\n\n".join(f"{str(self.reactions[i])}  {choice}" for i, choice in enumerate(choices, 1)),
+            description=f"**{title}**\n\n",
             timestamp=discord.utils.utcnow(),
             color=discord.colour.Color.gold(),
         )
-        embed.set_footer(text=f"Poll by {str(ctx.author)}")
-        msg = await ctx.send(embed=embed)
-        for i in range(1, len(choices) + 1):
-            await msg.add_reaction(self.reactions[i])
+        embed.set_footer(text=f"Poll by {str(interaction.user.display_name)}")
+        await interaction.response.send_message(embed=embed, ephemeral=True, view=PollButtons())
 
-    @poll.command()
-    async def show(self, ctx, message: str):
+    @app_commands.command()
+    async def show(self, interaction: core.InteractionType, message: str, ephemeral: bool = True):
         """Show a poll result"""
-        await ctx.message.delete()
-
         try:
             *_, channel_id, msg_id = message.split("/")
 
             try:
                 channel = self.__bot.get_channel(int(channel_id))
                 message = await channel.fetch_message(int(msg_id))
             except Exception:
-                return await ctx.send("Please provide the message ID/link for a valid poll")
+                return await interaction.response.send_message(
+                    "Please provide the message ID/link for a valid poll", ephemeral=True
+                )
         except Exception:
             try:
-                message = await ctx.channel.fetch_message(message)
+                message = await interaction.channel.fetch_message(int(message))
             except Exception:
-                return await ctx.send("Please provide the message ID/link for a valid poll")
+                return await interaction.response.send_message(
+                    "Please provide the message ID/link for a valid poll", ephemeral=True
+                )
 
         if self.poll_check(message):
             poll_embed = message.embeds[0]
@@ -142,10 +105,12 @@ async def show(self, ctx, message: str):
                 )
 
             embed.set_footer(text="Poll Result")
-            return await ctx.send(embed=embed)
+            return await interaction.response.send_message(embed=embed, ephemeral=ephemeral)
 
-        return await ctx.send("Please provide the message ID/link for a valid poll")
+        return await interaction.response.send_message(
+            "Please provide the message ID/link for a valid poll", ephemeral=True
+        )
 
 
-async def setup(bot):
-    await bot.add_cog(Polls(bot))
+async def setup(bot: commands.Bot):
+    await bot.add_cog(Polls(bot=bot))

bot/extensions/polls/events.py

@@ -0,0 +1,43 @@
+import discord
+from discord.ext import commands
+
+from bot import core
+
+
+class PollsEvents(commands.Cog):
+    """Events for polls in discord."""
+
+    def __init__(self, bot: core.DiscordBot):
+        self.bot = bot
+        self.emojis = ["1️⃣", "2️⃣", "3️⃣", "4️⃣", "5️⃣", "6️⃣", "7️⃣", "9️⃣", "🔟"]
+
+    def poll_check(self, message: discord.Message):
+        try:
+            embed = message.embeds[0]
+        except Exception:
+            return False
+        if str(embed.footer.text).count("Poll by") == 1:
+            return message.author == self.bot.user
+        return False
+
+    @commands.Cog.listener()
+    async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent):
+        channel: discord.TextChannel = self.bot.get_channel(payload.channel_id)
+        message: discord.Message = await channel.fetch_message(payload.message_id)
+
+        if payload.user_id == self.bot.user.id:
+            return
+
+        if not self.poll_check(message):
+            return
+
+        if str(payload.emoji) not in self.emojis:
+            return
+
+        for reaction in message.reactions:
+            if str(reaction) not in self.emojis:
+                return
+
+            if str(reaction.emoji) != str(payload.emoji):
+                user = self.bot.get_user(payload.user_id)
+                await message.remove_reaction(reaction.emoji, user)

bot/extensions/polls/views.py

@@ -0,0 +1,102 @@
+import discord
+from discord import ui
+
+
+class PollModal(ui.Modal):
+    name = ui.TextInput(label="Choice name", placeholder="Enter poll choice", max_length=50, required=True)
+
+    def __init__(self, var: discord.Interaction):
+        self.var = var
+        super().__init__(title="Poll options")
+
+    @property
+    def reactions(self):
+        return {
+            0: "1️⃣",
+            1: "2️⃣",
+            2: "3️⃣",
+            3: "4️⃣",
+            4: "5️⃣",
+            5: "6️⃣",
+            6: "7️⃣",
+            7: "8️⃣",
+            8: "9️⃣",
+            9: "🔟",
+        }
+
+    async def on_submit(self, interaction: discord.Interaction) -> None:
+        message = await self.var.followup.fetch_message(self.var.message.id)
+
+        # Determine the emoji to use based on the number of options
+        num = str(message.embeds[0].description).count("\n\n")
+        message.embeds[0].description += f"\n\n{str(self.reactions[num])}  {self.name}"
+
+        await message.edit(embed=message.embeds[0])
+        await interaction.response.defer()
+
+
+class PollButtons(ui.View):
+    def __init__(self, *, timeout=180):
+        super().__init__(timeout=timeout)
+
+    @property
+    def reactions(self):
+        return {
+            0: "1️⃣",
+            1: "2️⃣",
+            2: "3️⃣",
+            3: "4️⃣",
+            4: "5️⃣",
+            5: "6️⃣",
+            6: "7️⃣",
+            7: "8️⃣",
+            8: "9️⃣",
+            9: "🔟",
+        }
+
+    @discord.ui.button(label="Add Choice", style=discord.ButtonStyle.gray, emoji="➕")
+    async def add_choice(self, interaction: discord.Interaction, _button: ui.Button):
+        # Count the number of options
+        num = str(interaction.message.embeds[0].description).count("\n\n")
+        # If there are more than 10 options, return
+        if num >= 10:
+            return await interaction.response.send_message(
+                "You can't make a poll with more than 10 choices", ephemeral=True
+            )
+
+        modal = PollModal(var=interaction)
+        await interaction.response.send_modal(modal)
+
+    @discord.ui.button(label="Remove Choice", style=discord.ButtonStyle.gray, emoji="➖")
+    async def remove_choice(self, interaction: discord.Interaction, _button: ui.Button):
+        embed = interaction.message.embeds[0]
+
+        # If there are no options, return
+        if str(embed.description).count("\n\n") == 0:
+            return await interaction.response.send_message(
+                "You can't remove a choice from a poll with no choices", ephemeral=True
+            )
+
+        # Remove the last option
+        embed.description = "\n\n".join(embed.description.split("\n\n")[:-1])
+        await interaction.response.edit_message(embed=embed)
+
+    @discord.ui.button(label="Create Poll", style=discord.ButtonStyle.gray, emoji="📝")
+    async def create_poll(self, interaction: discord.Interaction, _button: ui.Button):
+        embed = interaction.message.embeds[0]
+
+        # If there are less than 2 options, return
+        if str(embed.description).count("\n\n") < 2:
+            return await interaction.response.send_message(
+                "You can't create a poll with less than 2 choices", ephemeral=True
+            )
+
+        message = await interaction.channel.send(embed=embed)
+
+        # Add reactions
+        for i in range(0, str(embed.description).count("\n\n")):
+            await message.add_reaction(self.reactions[i])
+
+        # Delete the original message
+        await interaction.response.defer()
+        await interaction.delete_original_response()

cli.py

@@ -137,10 +137,10 @@ async def main(ctx):
         "bot.extensions.tags",
         "bot.extensions.levelling",
         "bot.extensions.persistent_roles",
+        "bot.extensions.polls",
         "bot.cogs._help",
         "bot.cogs.clashofcode",
         "bot.cogs.roles",
-        "bot.cogs.poll",
     )
 
     intents = discord.Intents.all()