Added bot messages. Updated polls and self-roles.

Added bot messages, intended for moderators to be able to post as the bot.
These messages are also editable through a context menu.
Intended for rules/info messages etc.

Polls:
Added some error handling, as the bot only supports Unicode emojis.

Self-roles:
Edited the name of the context menu for updating a role-group message.

Heimdallr.py

@@ -134,4 +134,5 @@ if __name__ == "__main__":
     bot.load_extension("commands.infractions")
     bot.load_extension("commands.self_roles")
     bot.load_extension("commands.polls")
+    bot.load_extension("commands.bot_messages")
     bot.start(getenv("DISCORD_TOKEN"))

commands/bot_messages.py

@@ -0,0 +1,190 @@
+from copy import deepcopy
+import json
+from json.decoder import JSONDecodeError
+import logging
+from naff import (
+    Client,
+    Extension,
+    slash_command,
+    slash_option,
+    InteractionContext,
+    Modal,
+    ParagraphText,
+    ModalContext,
+    context_menu,
+    Permissions,
+    CommandTypes,
+    Message,
+)
+
+from database import BotMessages as BotMessagesModel
+
+
+message_creation_modal = Modal(
+    custom_id="bot-message-create",
+    title=f"Create a message as the bot",
+    components=[
+        ParagraphText(
+            custom_id="embeds",
+            label="Embeds",
+            placeholder="Embeds as JSON",
+            required=False,
+        ),
+        ParagraphText(
+            custom_id="content",
+            label="Text",
+            placeholder="Lorem ipsum dolor sit amet and so on",
+            required=False,
+        ),
+    ]
+)
+
+
+class BotMessages(Extension):
+    def __init__(self, client: Client) -> None:
+        self.client = client
+
+    @slash_command(
+        name="bot-message-create",
+        description="Create a message as the bot.",
+        dm_permission=False,
+        default_member_permissions=Permissions.MANAGE_GUILD,
+    )
+    async def bot_message_create_command(self, ctx: InteractionContext):
+        
+
+        await ctx.send_modal(message_creation_modal)
+        
+        modal_ctx: ModalContext = await self.client.wait_for_modal(message_creation_modal, author=ctx.author)
+        if modal_ctx.custom_id != "bot-message-create":
+            return
+        
+        embeds_string: str = modal_ctx.responses["embeds"]
+        content_string: str = modal_ctx.responses["content"]
+        if (
+            (embeds_string is None or embeds_string == "")
+            and (content_string is None or content_string == "")
+        ):
+            await modal_ctx.send(
+                "You must provide either an embed or text.",
+                ephemeral=True,
+            )
+            return
+
+        embed: dict | None = None
+        embeds: list | None = None
+        
+        try:
+            if embeds_string:
+                embeds_temp = json.loads(embeds_string)
+                if isinstance(embeds_temp, list):
+                    embeds = embeds_temp
+                elif isinstance(embeds_temp, dict):
+                    embed = embeds_temp
+        except JSONDecodeError:
+            await modal_ctx.send(
+                "The embeds were not valid JSON.",
+                ephemeral=True,
+            )
+            return
+        
+        msg = await ctx.channel.send(
+            content=content_string if content_string else None,
+            embed=embed,
+            embeds=embeds,
+        )
+        BotMessagesModel.create(
+            guild_id=msg.guild.id,
+            channel_id=msg.channel.id,
+            message_id=msg.id,
+        )
+        await modal_ctx.send(
+            "Message created!",
+            ephemeral=True,
+        )
+
+    @context_menu(
+        name="Edit bot message",
+        context_type=CommandTypes.MESSAGE,
+        dm_permission=False,
+        default_member_permissions=Permissions.MANAGE_GUILD,
+    )
+    async def edit_bot_message_context_menu(self, ctx: InteractionContext):
+        message: Message = ctx.target
+        if message.author.id != self.client.user.id:
+            await ctx.send(
+                "This is not a bot message.",
+                ephemeral=True,
+            )
+            return
+
+        bot_message: BotMessagesModel | None = BotMessagesModel.get(
+            BotMessagesModel.guild_id == message.channel.guild.id,
+            BotMessagesModel.channel_id == message.channel.id,
+            BotMessagesModel.message_id == message.id,
+        )
+
+        if bot_message is None:
+            await ctx.send(
+                "This is not an editable bot message.",
+                ephemeral=True,
+            )
+            return
+        modal = deepcopy(message_creation_modal)
+        modal.title = "Edit bot message"
+        modal.components[0].value = json.dumps(
+            [e.to_dict() for e in message.embeds] if message.embeds else "",
+            indent=4,
+        )
+        modal.components[1].value = message.content
+
+        await ctx.send_modal(modal)
+        
+        modal_ctx: ModalContext = await self.client.wait_for_modal(modal, author=ctx.author)
+        if modal_ctx.custom_id != "bot-message-create":
+            return
+        
+        embeds_string: str = modal_ctx.responses["embeds"]
+        content_string: str = modal_ctx.responses["content"]
+        if (
+            (embeds_string is None or embeds_string == "")
+            and (content_string is None or content_string == "")
+        ):
+            await modal_ctx.send(
+                "You must provide either an embed or text.",
+                ephemeral=True,
+            )
+            return
+
+        embed: dict | None = None
+        embeds: list | None = None
+        
+        try:
+            if embeds_string:
+                embeds_temp = json.loads(embeds_string)
+                if isinstance(embeds_temp, list):
+                    embeds = embeds_temp
+                elif isinstance(embeds_temp, dict):
+                    embed = embeds_temp
+        except JSONDecodeError:
+            await modal_ctx.send(
+                "The embeds were not valid JSON.",
+                ephemeral=True,
+            )
+            return
+        
+        await message.edit(
+            content=content_string if content_string else None,
+            embed=embed,
+            embeds=embeds,
+        )
+        await modal_ctx.send(
+            "Message edited!",
+            ephemeral=True,
+        )
+
+
+def setup(client: Client):
+    BotMessagesModel.create_table()
+    BotMessages(client)
+    logging.info("BotMessages extension loaded")

commands/polls.py

@@ -23,8 +23,10 @@ from naff import (
     Task,
     IntervalTrigger,
     GuildText,
+    PartialEmoji,
 )
 from naff.api import events
+from naff.client.errors import HTTPException
 
 from database import Polls as PollsModel, PollVotes as PollVotesModel
 from peewee import fn
@@ -70,7 +72,7 @@ def generate_poll_embed(
     embed = Embed(
         title=title,
         description=("\n".join(data) if data else None),
-        )
+    )
     sum_votes = sum(votes)
 
     for i, (emoji, option) in enumerate(options):
@@ -217,9 +219,13 @@ class Polls(Extension):
 
         buttons: List[Button] = []
         for i, option in enumerate(options):
+            try:
+                emoji = PartialEmoji.from_str((option[0] or num_to_emoji(i + 1)))
+            except ValueError:
+                emoji = num_to_emoji(i + 1)
             buttons.append(
                 Button(
-                    emoji=(option[0] or num_to_emoji(i + 1)),
+                    emoji=emoji,
                     style=ButtonStyles.PRIMARY,
                     custom_id=f"poll-vote:{poll_entry.id}:{i}",
                 )
@@ -227,9 +233,10 @@ class Polls(Extension):
 
         buttons.append(
             Button(
-                label="Delete",
+                emoji="🔒",
+                label="Lock",
                 style=ButtonStyles.DANGER,
-                custom_id=f"poll-delete:{poll_entry.id}",
+                custom_id=f"poll-lock:{poll_entry.id}",
             )
         )
         embed = generate_poll_embed(
@@ -239,10 +246,18 @@ class Polls(Extension):
             multiple_choice=multiple_choice,
             expires=duration,
         )
-        poll_message: Message = await modal_ctx.send(
+        try:
-            embed=embed,
+            poll_message: Message = await modal_ctx.send(
-            components=spread_to_rows(*buttons),
+                embed=embed,
-        )
+                components=spread_to_rows(*buttons),
+            )
+        except HTTPException as e:
+            logging.error(f"Error sending poll: {e}")
+            await modal_ctx.send(
+                "Error during poll creation. NB: You can not use server-specific emojis",
+                ephemeral=True,
+            )
+            return
 
         poll_entry.message_id = poll_message.id
         poll_entry.channel_id = poll_message.channel.id
@@ -347,7 +362,7 @@ class Polls(Extension):
 
             await ctx.message.edit(embed=embed)
 
-        elif ctx.custom_id.startswith("poll-delete:"):
+        elif ctx.custom_id.startswith("poll-lock:"):
             poll_id = ctx.custom_id.split(":", 1)[1]
 
             poll_entry: PollsModel | None = PollsModel.get_or_none(
@@ -357,17 +372,17 @@ class Polls(Extension):
                 await ctx.send("That poll doesn't exist.")
                 return
 
-            if (
+            if not ctx.author.id == int(
-                not ctx.author.id == int(poll_entry.author_id)
+                poll_entry.author_id
-                or not ctx.author.has_permission(Permissions.MANAGE_MESSAGES)
+            ) or not ctx.author.has_permission(Permissions.MANAGE_MESSAGES):
-            ):
+                await ctx.send("You don't have permission to lock that poll.")
-                await ctx.send("You don't have permission to delete that poll.")
                 return
 
-            poll_entry.delete_instance()
+            poll_entry.expires = datetime.now() - timedelta(minutes=1)
-            PollVotesModel.delete().where(PollVotesModel.poll_id == poll_id).execute()
+            poll_entry.save()
-            await ctx.message.delete()
+
-            await ctx.send("Poll deleted.")
+            await self.poll_expiry_check()
+            await ctx.send("Poll locked.")
 
     @Task.create(IntervalTrigger(minutes=1))
     async def poll_expiry_check(self):
@@ -376,7 +391,9 @@ class Polls(Extension):
 
         polls_q: List[PollsModel] = PollsModel.select().where(PollsModel.expires < now)
         for poll_entry in polls_q:
-            channel: GuildText = await self.client.fetch_channel(int(poll_entry.channel_id))
+            channel: GuildText = await self.client.fetch_channel(
+                int(poll_entry.channel_id)
+            )
             if not channel:
                 continue
 
@@ -387,8 +404,6 @@ class Polls(Extension):
             await message.edit(components=[])
 
 
-
-
 def setup(client: Client):
     PollsModel.create_table()
     PollVotesModel.create_table()

commands/self_roles.py

@@ -536,7 +536,7 @@ class SelfRoles(Extension):
             )
 
     @context_menu(
-        name="Regenerate group",
+        name="Regenerate role group message",
         context_type=CommandTypes.MESSAGE,
         dm_permission=False,
         default_member_permissions=Permissions.MANAGE_ROLES,