From 0c382986cfae955d8930876a34529e363a33d77f Mon Sep 17 00:00:00 2001 From: Zevaryx Date: Thu, 3 Feb 2022 15:46:17 -0700 Subject: [PATCH] Migrate rolegiver, use components for add/remove roles, closes #101 --- jarvis/cogs/rolegiver.py | 340 +++++++++++++++++++++------------------ 1 file changed, 187 insertions(+), 153 deletions(-) diff --git a/jarvis/cogs/rolegiver.py b/jarvis/cogs/rolegiver.py index 4047ac0..fbeaa4a 100644 --- a/jarvis/cogs/rolegiver.py +++ b/jarvis/cogs/rolegiver.py @@ -1,44 +1,38 @@ """J.A.R.V.I.S. Role Giver Cog.""" import asyncio -from discord import Role -from discord.ext import commands -from discord_slash import SlashContext, cog_ext -from discord_slash.utils.manage_commands import create_option -from discord_slash.utils.manage_components import ( - create_actionrow, - create_select, - create_select_option, - wait_for_component, +from dis_snek import InteractionContext, Permissions, Scale, Snake +from dis_snek.client.utils import get +from dis_snek.models.discord.components import ActionRow, Select, SelectOption +from dis_snek.models.discord.embed import EmbedField +from dis_snek.models.discord.role import Role +from dis_snek.models.snek.application_commands import ( + OptionTypes, + slash_command, + slash_option, ) +from dis_snek.models.snek.command import cooldown +from dis_snek.models.snek.cooldowns import Buckets from jarvis.db.models import Rolegiver from jarvis.utils import build_embed -from jarvis.utils.field import Field from jarvis.utils.permissions import admin_or_permissions -class RolegiverCog(commands.Cog): +class RolegiverCog(Scale): """J.A.R.V.I.S. Role Giver Cog.""" - def __init__(self, bot: commands.Bot): + def __init__(self, bot: Snake): self.bot = bot - @cog_ext.cog_subcommand( - base="rolegiver", - name="add", - description="Add a role to rolegiver", - options=[ - create_option( - name="role", - description="Role to add", - option_type=8, - required=True, - ) - ], + @slash_command( + name="rolegiver", sub_cmd_name="add", sub_cmd_description="Add a role to rolegiver" ) - @admin_or_permissions(manage_guild=True) - async def _rolegiver_add(self, ctx: SlashContext, role: Role) -> None: + @slash_option( + name="role", description="Role to add", optin_type=OptionTypes.ROLE, required=True + ) + @admin_or_permissions(Permissions.MANAGE_GUILD) + async def _rolegiver_add(self, ctx: InteractionContext, role: Role) -> None: setting = Rolegiver.objects(guild=ctx.guild.id).first() if setting and role.id in setting.roles: await ctx.send("Role already in rolegiver", hidden=True) @@ -58,7 +52,7 @@ class RolegiverCog(commands.Cog): for role_id in setting.roles: if role_id == role.id: continue - e_role = ctx.guild.get_role(role_id) + e_role = await ctx.guild.get_role(role_id) if not e_role: continue roles.append(e_role) @@ -67,8 +61,8 @@ class RolegiverCog(commands.Cog): value = "\n".join([r.mention for r in roles]) if roles else "None" fields = [ - Field(name="New Role", value=f"{role.mention}"), - Field(name="Existing Role(s)", value=value), + EmbedField(name="New Role", value=f"{role.mention}"), + EmbedField(name="Existing Role(s)", value=value), ] embed = build_embed( @@ -79,21 +73,19 @@ class RolegiverCog(commands.Cog): embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_author( - name=ctx.author.nick if ctx.author.nick else ctx.author.name, - icon_url=ctx.author.avatar_url, + name=ctx.author.display_name, + icon_url=ctx.author.display_avatar, ) - embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") + embed.set_footer(text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}") await ctx.send(embed=embed) - @cog_ext.cog_subcommand( - base="rolegiver", - name="remove", - description="Remove a role from rolegiver", + @slash_command( + name="rolegiver", sub_cmd_name="remove", sub_cmd_description="Remove a role from rolegiver" ) - @admin_or_permissions(manage_guild=True) - async def _rolegiver_remove(self, ctx: SlashContext) -> None: + @admin_or_permissions(Permissions.MANAGE_GUILD) + async def _rolegiver_remove(self, ctx: InteractionContext) -> None: setting = Rolegiver.objects(guild=ctx.guild.id).first() if not setting or (setting and not setting.roles): await ctx.send("Rolegiver has no roles", hidden=True) @@ -101,28 +93,28 @@ class RolegiverCog(commands.Cog): options = [] for role in setting.roles: - role: Role = ctx.guild.get_role(role) - option = create_select_option(label=role.name, value=str(role.id)) + role: Role = await ctx.guild.get_role(role) + option = SelectOption(label=role.name, value=str(role.id)) options.append(option) - select = create_select( + select = Select( options=options, custom_id="to_delete", placeholder="Select roles to remove", min_values=1, max_values=len(options), ) - components = [create_actionrow(select)] + components = [ActionRow(select)] message = await ctx.send(content="\u200b", components=components) try: - context = await wait_for_component( + context = await self.bot.wait_for_component( self.bot, check=lambda x: ctx.author.id == x.author.id, message=message, timeout=60 * 1, ) - for to_delete in context.selected_options: + for to_delete in context.context.values: setting.roles.remove(int(to_delete)) setting.save() for row in components: @@ -131,7 +123,7 @@ class RolegiverCog(commands.Cog): roles = [] for role_id in setting.roles: - e_role = ctx.guild.get_role(role_id) + e_role = await ctx.guild.get_role(role_id) if not e_role: continue roles.append(e_role) @@ -141,8 +133,8 @@ class RolegiverCog(commands.Cog): value = "\n".join([r.mention for r in roles]) if roles else "None" fields = [ - Field(name="Removed Role", value=f"{role.mention}"), - Field(name="Remaining Role(s)", value=value), + EmbedField(name="Removed Role", value=f"{role.mention}"), + EmbedField(name="Remaining Role(s)", value=value), ] embed = build_embed( @@ -153,13 +145,15 @@ class RolegiverCog(commands.Cog): embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_author( - name=ctx.author.nick if ctx.author.nick else ctx.author.name, - icon_url=ctx.author.avatar_url, + name=ctx.author.display_name, + icon_url=ctx.author.display_avatar, ) - embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") + embed.set_footer( + text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}" + ) - await context.edit_origin( + await context.context.edit_origin( content=f"Removed {len(context.selected_options)} role(s)", embed=embed, components=components, @@ -170,12 +164,8 @@ class RolegiverCog(commands.Cog): component["disabled"] = True await message.edit(components=components) - @cog_ext.cog_subcommand( - base="rolegiver", - name="list", - description="List roles rolegiver", - ) - async def _rolegiver_list(self, ctx: SlashContext) -> None: + @slash_command(name="rolegiver", sub_cmd_name="list", description="List rolegiver roles") + async def _rolegiver_list(self, ctx: InteractionContext) -> None: setting = Rolegiver.objects(guild=ctx.guild.id).first() if not setting or (setting and not setting.roles): await ctx.send("Rolegiver has no roles", hidden=True) @@ -183,7 +173,7 @@ class RolegiverCog(commands.Cog): roles = [] for role_id in setting.roles: - e_role = ctx.guild.get_role(role_id) + e_role = await ctx.guild.get_role(role_id) if not e_role: continue roles.append(e_role) @@ -201,21 +191,17 @@ class RolegiverCog(commands.Cog): embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_author( - name=ctx.author.nick if ctx.author.nick else ctx.author.name, - icon_url=ctx.author.avatar_url, + name=ctx.author.nick if ctx.author.nick else ctx.author.username, + icon_url=ctx.author.display_avatar, ) - embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") + embed.set_footer(text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}") await ctx.send(embed=embed) - @cog_ext.cog_subcommand( - base="role", - name="get", - description="Get a role from rolegiver", - ) - @commands.cooldown(1, 10, commands.BucketType.user) - async def _role_get(self, ctx: SlashContext) -> None: + @slash_command(name="role", sub_cmd_name="get", sub_cmd_description="Get a role") + @cooldown(bucket=Buckets.USER, rate=1, interval=10) + async def _role_get(self, ctx: InteractionContext) -> None: setting = Rolegiver.objects(guild=ctx.guild.id).first() if not setting or (setting and not setting.roles): await ctx.send("Rolegiver has no roles", hidden=True) @@ -223,117 +209,164 @@ class RolegiverCog(commands.Cog): options = [] for role in setting.roles: - role: Role = ctx.guild.get_role(role) - option = create_select_option(label=role.name, value=str(role.id)) + role: Role = await ctx.guild.get_role(role) + option = SelectOption(label=role.name, value=str(role.id)) options.append(option) - select = create_select( + select = Select( options=options, custom_id="to_delete", placeholder="Select roles to remove", min_values=1, max_values=len(options), ) - components = [create_actionrow(select)] + components = [ActionRow(select)] - _ = await ctx.send(content="\u200b", components=components) + message = await ctx.send(content="\u200b", components=components) - await ctx.author.add_roles(role, reason="Rolegiver") - - roles = ctx.author.roles - if roles: - roles.sort(key=lambda x: -x.position) - _ = roles.pop(-1) - - value = "\n".join([r.mention for r in roles]) if roles else "None" - fields = [ - Field(name="Added Role", value=f"{role.mention}"), - Field(name="Prior Role(s)", value=value), - ] - - embed = build_embed( - title="User Given Role", - description=f"{role.mention} given to {ctx.author.mention}", - fields=fields, - ) - - embed.set_thumbnail(url=ctx.guild.icon_url) - embed.set_author( - name=ctx.author.nick if ctx.author.nick else ctx.author.name, - icon_url=ctx.author.avatar_url, - ) - - embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") - - await ctx.send(embed=embed) - - @cog_ext.cog_subcommand( - base="role", - name="forfeit", - description="Have rolegiver take away role", - options=[ - create_option( - name="role", - description="Role to remove", - option_type=8, - required=True, + try: + context = await self.bot.wait_for_component( + check=lambda x: ctx.author.id == x.author_id, + messages=message, + timeout=60 * 5, ) - ], - ) - @commands.cooldown(1, 10, commands.BucketType.user) - async def _role_forfeit(self, ctx: SlashContext, role: Role) -> None: + + added_roles = [] + for role in context.context.values: + role = await ctx.guild.get_role(int(role)) + added_roles.append(role) + await ctx.author.add_role(role, reason="Rolegiver") + + roles = ctx.author.roles + if roles: + roles.sort(key=lambda x: -x.position) + _ = roles.pop(-1) + + avalue = "\n".join([r.mention for r in added_roles]) if added_roles else "None" + value = "\n".join([r.mention for r in roles]) if roles else "None" + fields = [ + EmbedField(name="Added Role(s)", value=avalue), + EmbedField(name="Prior Role(s)", value=value), + ] + + embed = build_embed( + title="User Given Role", + description=f"{role.mention} given to {ctx.author.mention}", + fields=fields, + ) + + embed.set_thumbnail(url=ctx.guild.icon_url) + embed.set_author( + name=ctx.author.display_name, + icon_url=ctx.author.display_avatar, + ) + + embed.set_footer( + text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}" + ) + + for row in components: + for component in row["components"]: + component["disabled"] = True + + await context.context.edit_origin(embed=embed, content="\u200b", components=components) + except asyncio.TimeoutError: + for row in components: + for component in row["components"]: + component["disabled"] = True + await message.edit(components=components) + + @slash_command(name="role", sub_cmd_name="remove", sub_cmd_description="Remove a role") + @cooldown(bucket=Buckets.USER, rate=1, interval=10) + async def _role_remove(self, ctx: InteractionContext) -> None: + user_roles = ctx.author.roles + setting = Rolegiver.objects(guild=ctx.guild.id).first() if not setting or (setting and not setting.roles): await ctx.send("Rolegiver has no roles", hidden=True) return - elif role.id not in setting.roles: - await ctx.send("Role not in rolegiver", hidden=True) - return - elif role not in ctx.author.roles: - await ctx.send("You do not have that role", hidden=True) + elif not any(x.id in setting.roles for x in user_roles): + await ctx.send("You have no rolegiver roles", hidden=True) return - await ctx.author.remove_roles(role, reason="Rolegiver") + valid = list(filter(lambda x: x.id in setting.roles, user_roles)) + options = [] + for role in valid: + option = SelectOption(label=role.name, value=str(role.id)) + options.append(option) - roles = ctx.author.roles - if roles: - roles.sort(key=lambda x: -x.position) - _ = roles.pop(-1) - - value = "\n".join([r.mention for r in roles]) if roles else "None" - fields = [ - Field(name="Taken Role", value=f"{role.mention}"), - Field(name="Remaining Role(s)", value=value), - ] - - embed = build_embed( - title="User Forfeited Role", - description=f"{role.mention} taken from {ctx.author.mention}", - fields=fields, + select = Select( + options=options, + custom_id="to_remove", + placeholder="Select roles to remove", + min_values=1, + max_values=len(options), ) + components = [ActionRow(select)] - embed.set_thumbnail(url=ctx.guild.icon_url) - embed.set_author( - name=ctx.author.nick if ctx.author.nick else ctx.author.name, - icon_url=ctx.author.avatar_url, - ) + message = await ctx.send(content="\u200b", components=components) - embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") + try: + context = await self.bot.wait_for_component( + check=lambda x: ctx.author.id == x.author_id, + messages=message, + timeout=60 * 5, + ) - await ctx.send(embed=embed) + removed_roles = [] + for to_remove in context.context.values: + role = get(user_roles, id=int(to_remove)) + await ctx.author.remove_role(role, reason="Rolegiver") + user_roles.remove(role) + removed_roles.append(role) - @cog_ext.cog_subcommand( - base="rolegiver", - name="cleanup", - description="Cleanup rolegiver roles", + user_roles.sort(key=lambda x: -x.position) + _ = user_roles.pop(-1) + + value = "\n".join([r.mention for r in user_roles]) if user_roles else "None" + rvalue = "\n".join([r.mention for r in removed_roles]) if removed_roles else "None" + fields = [ + EmbedField(name="Removed Role(s)", value=rvalue), + EmbedField(name="Remaining Role(s)", value=value), + ] + + embed = build_embed( + title="User Forfeited Role", + description=f"{role.mention} taken from {ctx.author.mention}", + fields=fields, + ) + + embed.set_thumbnail(url=ctx.guild.icon_url) + embed.set_author( + name=ctx.author.display_name, + icon_url=ctx.author.display_avatar, + ) + + embed.set_footer( + text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}" + ) + + for row in components: + for component in row["components"]: + component["disabled"] = True + + await context.context.edit_origin(embed=embed, components=components, content="\u200b") + + except asyncio.TimeoutError: + for row in components: + for component in row["components"]: + component["disabled"] = True + await message.edit(components=components) + + @slash_command( + name="rolegiver", sub_cmd_name="cleanup", description="Removed deleted roles from rolegiver" ) - @admin_or_permissions(manage_guild=True) - async def _rolegiver_cleanup(self, ctx: SlashContext) -> None: + @admin_or_permissions(Permissions.MANAGE_GUILD) + async def _rolegiver_cleanup(self, ctx: InteractionContext) -> None: setting = Rolegiver.objects(guild=ctx.guild.id).first() if not setting or not setting.roles: await ctx.send("Rolegiver has no roles", hidden=True) - guild_roles = await ctx.guild.fetch_roles() - guild_role_ids = [x.id for x in guild_roles] + guild_role_ids = [r.id for r in ctx.guild.roles] for role_id in setting.roles: if role_id not in guild_role_ids: setting.roles.remove(role_id) @@ -342,6 +375,7 @@ class RolegiverCog(commands.Cog): await ctx.send("Rolegiver cleanup finished") -def setup(bot: commands.Bot) -> None: +def setup(bot: Snake) -> None: """Add RolegiverCog to J.A.R.V.I.S.""" bot.add_cog(RolegiverCog(bot)) + bot.add_cog(RolegiverCog(bot))