From 43cd7ebe989d81b521a6f8acb06f1e7ccdc251a4 Mon Sep 17 00:00:00 2001 From: Sylte Date: Mon, 6 Nov 2023 22:56:24 +0100 Subject: [PATCH] Update myrole command to use existing functions, update embeds and event code. --- bot/extensions/custom_roles/commands.py | 87 +++++++++++++------------ bot/extensions/custom_roles/events.py | 56 +++++++++------- bot/models/custom_roles.py | 10 +-- 3 files changed, 82 insertions(+), 71 deletions(-) diff --git a/bot/extensions/custom_roles/commands.py b/bot/extensions/custom_roles/commands.py index bd630c7e..194ea1b5 100644 --- a/bot/extensions/custom_roles/commands.py +++ b/bot/extensions/custom_roles/commands.py @@ -1,9 +1,8 @@ import datetime import discord -from discord import app_commands +from discord import app_commands, utils from discord.ext import commands -from discord.utils import escape_markdown, format_dt from bot import core from bot.models import CustomRole @@ -13,64 +12,65 @@ class CustomRoles(commands.Cog): def __init__(self, bot): self.bot = bot - def role_embed(self, heading: str, user: discord.Member, role: discord.Role): + self.color_converter = commands.ColorConverter() + + @staticmethod + def role_embed(heading: str, user: discord.Member, role: discord.Role): embed = discord.Embed( - description=f"### {heading} {user.mention}\n\n**Name**\n{escape_markdown(role.name)}\n**Color**\n" - f"{role.color}\n**Created**\n {format_dt(role.created_at)}", + title=heading, color=role.color, timestamp=datetime.datetime.utcnow(), ) - embed.set_footer(text=user.name) + embed.add_field(name="Name", value=utils.escape_markdown(role.name)) + embed.add_field(name="Color", value=str(role.color)) + embed.add_field(name="Created at", value=utils.format_dt(role.created_at)) embed.set_thumbnail(url=user.avatar) return embed - @app_commands.command(description="Manage custom role") + @app_commands.command() @app_commands.default_permissions(administrator=True) - @app_commands.describe(name="Updating Custom Role name", color="Updating custom role name") + @app_commands.describe(name="New name", color="New color") async def myrole(self, interaction: core.InteractionType, name: str = None, color: str = None): - if color: + """Manage your custom role""" + if color is not None: try: - color = discord.Color(int(color, 16)) - except ValueError: - return await interaction.response.send_message("Invalid HEX value", ephemeral=True) + color = await self.color_converter.convert(None, color) # noqa + except commands.BadColourArgument as e: + return await interaction.response.send_message(str(e), ephemeral=True) - query = """SELECT * FROM custom_roles - WHERE guild_id = $1 AND user_id = $2""" + query = "SELECT * FROM custom_roles WHERE guild_id = $1 AND user_id = $2" before = await CustomRole.fetchrow(query, interaction.guild.id, interaction.user.id) - # Insert data if it doesn't exist if before is None: - # Validate the name parameter if name is None: return await interaction.response.send_message("You don't have a custom role yet!", ephemeral=True) # Create and assign the role to user - discord_role = await interaction.guild.create_role(name=name, colour=color or discord.Color.default()) - await interaction.user.add_roles(discord_role) - - query = """INSERT INTO custom_roles (user_id, guild_id, role_id, name, color) - VALUES ($1, $2, $3, $4, $5) - RETURNING *""" - record = await CustomRole.fetchrow( - query, - interaction.user.id, - interaction.guild.id, - discord_role.id, - discord_role.name, - str(discord_role.color), + role = await interaction.guild.create_role(name=name, colour=color or discord.Color.random()) + + record = await CustomRole.ensure_exists( + guild_id=interaction.guild.id, + user_id=interaction.user.id, + role_id=role.id, + name=role.name, + color=role.color.value, ) - # Persist the role + + self.bot.dispatch("custom_role_create", custom_role=record) self.bot.dispatch( - "persist_roles", guild_id=interaction.guild.id, user_id=interaction.user.id, role_ids=[discord_role.id] + "persist_roles", + guild_id=interaction.guild.id, + user_id=interaction.user.id, + role_ids=[role.id], ) - # Log the role - self.bot.dispatch("custom_role_create", data=record) return await interaction.response.send_message( - embed=self.role_embed("**Custom Role has been assigned**", interaction.user, discord_role), + embed=self.role_embed("**Custom Role has been assigned**", interaction.user, role), ephemeral=True, ) + role = interaction.guild.get_role(before.role_id) + # Return role information if no parameter is passed if not name and not color: return await interaction.response.send_message( @@ -78,18 +78,19 @@ async def myrole(self, interaction: core.InteractionType, name: str = None, colo ephemeral=True, ) - # Editing the role - await interaction.guild.get_role(before.role_id).edit( + await role.edit( name=name or before.name, - colour=color or discord.Color(int(before.color.lstrip("#"), 16)), + colour=color or discord.Color(int(before.color)), + ) + + after = await CustomRole.ensure_exists( + guild_id=interaction.guild.id, + user_id=interaction.user.id, + role_id=role.id, + name=role.name, + color=role.color.value, ) - # Update data in DB - query = """UPDATE custom_roles - SET name = $2, color = $3 - WHERE role_id = $1 - RETURNING *""" - after = await CustomRole.fetchrow(query, before.role_id, name, color or before.color) self.bot.dispatch("custom_role_update", before, after) return await interaction.response.send_message( diff --git a/bot/extensions/custom_roles/events.py b/bot/extensions/custom_roles/events.py index 4880a925..aaad6302 100644 --- a/bot/extensions/custom_roles/events.py +++ b/bot/extensions/custom_roles/events.py @@ -1,4 +1,5 @@ import datetime +import time import discord from discord.ext import commands @@ -21,38 +22,49 @@ def custom_roles_logs_channel(self) -> discord.TextChannel | None: @commands.Cog.listener() async def on_guild_role_update(self, before: discord.Role, after: discord.Role): - try: - if datetime.datetime.now() - self.updated_at[before.id] < datetime.timedelta(0, 10): - query = """UPDATE custom_roles - SET name = $2, color = $3 - WHERE role_id = $1 - RETURNING *""" - await CustomRole.fetchrow(query, before.id, after.name, after.color) + last_update = self.updated_at.get(before.id) - except KeyError: - pass + # Ignore events less than 10 seconds since a user updated their role + if last_update is not None: + if time.time() - last_update < 10: + return + + query = """ + UPDATE custom_roles + SET name = $2, color = $3 + WHERE role_id = $1 + """ + await CustomRole.execute(query, after.id, after.name, after.color.value) @commands.Cog.listener() - async def on_custom_role_create(self, data: CustomRole): + async def on_custom_role_create(self, custom_role: CustomRole): """Logs the creation of new role""" + self.updated_at[custom_role.role_id] = time.time() + + user = await self.bot.fetch_user(custom_role.user_id) + embed = discord.Embed( - description=f"### <@{data.user_id}> Custom Role Created", + title="Custom Role Created", color=discord.Color.brand_green(), timestamp=datetime.datetime.utcnow(), ) - embed.add_field(name="Name", value=data.name) - embed.add_field(name="Color", value=str(data.color)) - embed.set_thumbnail(url=(await self.bot.fetch_user(data.user_id)).avatar) - - self.updated_at.setdefault(data.role_id, datetime.datetime.now()) + embed.set_author(name=user.display_name, icon_url=user.display_avatar) + embed.add_field(name="Name", value=custom_role.name) + embed.add_field(name="Color", value="#" + hex(custom_role.color)[2:]) + embed.set_thumbnail(url=user.avatar) + embed.set_footer(text=f"user_id: {custom_role.user_id}") return await self.custom_roles_logs_channel.send(embed=embed) @commands.Cog.listener() async def on_custom_role_update(self, before: CustomRole, after: CustomRole): """Logs the update of custom role.""" + self.updated_at[after.role_id] = time.time() + + user = await self.bot.fetch_user(after.user_id) + embed = discord.Embed( - description=f"### <@{before.user_id}> Custom Role Updated", + title="Custom Role Updated", color=discord.Color.brand_green(), timestamp=datetime.datetime.utcnow(), ) @@ -60,11 +72,9 @@ async def on_custom_role_update(self, before: CustomRole, after: CustomRole): embed.add_field(name="Old Name", value=before.name) embed.add_field(name="New Name", value=after.name) embed.add_field(name="\u200B", value="\u200B") - embed.add_field(name="Old Color", value=before.color) - embed.add_field(name="New Color", value=str(after.color)) - embed.add_field(name="\u200B", value="\u200B") - embed.set_thumbnail(url=(await self.bot.fetch_user(before.user_id)).avatar) - - self.updated_at.setdefault(after.role_id, datetime.datetime.now()) + embed.add_field(name="Old Color", value="#" + hex(before.color)[2:]) + embed.add_field(name="New Color", value="#" + hex(after.color)[2:]) + embed.set_thumbnail(url=user.avatar) + embed.set_footer(text=f"user_id: {after.user_id}") return await self.custom_roles_logs_channel.send(embed=embed) diff --git a/bot/models/custom_roles.py b/bot/models/custom_roles.py index 1a51dfc2..35736fae 100644 --- a/bot/models/custom_roles.py +++ b/bot/models/custom_roles.py @@ -7,14 +7,14 @@ class CustomRole(Model): guild_id: int role_id: int name: str - color: str + color: int @classmethod - async def ensure_exists(cls, guild_id: int, role_id: int, name: str, color: str): + async def ensure_exists(cls, guild_id: int, role_id: int, name: str, color: int, user_id: int = None): """Inserts or updates the custom role.""" query = """ - INSERT INTO custom_roles (guild_id, role_id, name, color) - VALUES ($1, $2, $3, $4) + INSERT INTO custom_roles (guild_id, role_id, name, color, user_id) + VALUES ($1, $2, $3, $4, $5) ON CONFLICT (guild_id, role_id) DO UPDATE SET name = $3, @@ -22,4 +22,4 @@ async def ensure_exists(cls, guild_id: int, role_id: int, name: str, color: str) RETURNING * """ - return await cls.fetchrow(query, guild_id, role_id, name, color) + return await cls.fetchrow(query, guild_id, role_id, name, color, user_id)