From 9bb5f2def8a8e3bd6228c5adee7c3697a46b27d1 Mon Sep 17 00:00:00 2001 From: Zevaryx Date: Sun, 16 Oct 2022 18:04:53 -0600 Subject: [PATCH] Formatting, change errors to be deferred --- jarvis/client/errors.py | 1 + jarvis/cogs/rolegiver.py | 167 ++++++++++++++++++++++----------------- jarvis/cogs/util.py | 7 ++ 3 files changed, 102 insertions(+), 73 deletions(-) diff --git a/jarvis/client/errors.py b/jarvis/client/errors.py index 5d15170..64a390c 100644 --- a/jarvis/client/errors.py +++ b/jarvis/client/errors.py @@ -102,6 +102,7 @@ class ErrorMixin: ) await ctx.send("Whoops! Encountered an error. The error has been logged.", ephemeral=True) try: + await ctx.defer(ephemeral=True) return await super().on_command_error(ctx, error, *args, **kwargs) except Exception as e: self.logger.error("Uncaught exception", exc_info=e) diff --git a/jarvis/cogs/rolegiver.py b/jarvis/cogs/rolegiver.py index b4727c8..4d5edb5 100644 --- a/jarvis/cogs/rolegiver.py +++ b/jarvis/cogs/rolegiver.py @@ -1,10 +1,18 @@ """JARVIS Role Giver Cog.""" import asyncio import logging +from typing import Dict from jarvis_core.db import q from jarvis_core.db.models import Rolegiver -from naff import Client, Extension, InteractionContext, Permissions +from naff import ( + AutocompleteContext, + Client, + Extension, + InteractionContext, + Permissions, + listen, +) from naff.client.utils.misc_utils import get from naff.models.discord.components import ActionRow, Button, Select, SelectOption from naff.models.discord.embed import EmbedField @@ -17,6 +25,7 @@ from naff.models.naff.application_commands import ( ) from naff.models.naff.command import check, cooldown from naff.models.naff.cooldowns import Buckets +from thefuzz import process from jarvis.utils import build_embed from jarvis.utils.permissions import admin_or_permissions @@ -28,6 +37,24 @@ class RolegiverCog(Extension): def __init__(self, bot: Client): self.bot = bot self.logger = logging.getLogger(__name__) + self.cache: Dict[int, Dict[str, int]] = {} + + @listen() + async def on_ready(self) -> None: + """NAFF on_ready hook for loading cache.""" + all_rolegivers = await Rolegiver.find({}).to_list(None) + for rolegiver in all_rolegivers: + guild = await self.bot.fetch_guild(rolegiver.guild) + if not guild: + await rolegiver.delete() + continue + role = await guild.fetch_role(rolegiver.role) + if not role: + await rolegiver.delete() + continue + if guild.id not in self.cache: + self.cache[guild.id] = {} + self.cache[guild.id][role.name] = role.id rolegiver = SlashCommand(name="rolegiver", description="Allow users to choose their own roles") @@ -59,10 +86,6 @@ class RolegiverCog(Extension): setting.roles = setting.roles or [] - if len(setting.roles) >= 20: - await ctx.send("You can only have 20 roles in the rolegiver", ephemeral=True) - return - setting.roles.append(role.id) await setting.commit() @@ -97,87 +120,68 @@ class RolegiverCog(Extension): ) await ctx.send(embeds=embed, components=components) + if ctx.guild.id not in self.cache: + self.cache[ctx.guild.id] = {} + self.cache[ctx.guild.id][role.name] = role.id + @rolegiver.subcommand(sub_cmd_name="remove", sub_cmd_description="Remove a role from rolegiver") + @slash_option( + name="role", + description="Name of role to add", + opt_type=OptionTypes.STRING, + required=True, + autocomplete=True, + ) @check(admin_or_permissions(Permissions.MANAGE_GUILD)) - async def _rolegiver_remove(self, ctx: InteractionContext) -> None: + async def _rolegiver_remove(self, ctx: InteractionContext, role: str) -> None: setting = await Rolegiver.find_one(q(guild=ctx.guild.id)) if not setting or (setting and not setting.roles): await ctx.send("Rolegiver has no roles", ephemeral=True) return - options = [] - for role in setting.roles: - role: Role = await ctx.guild.fetch_role(role) - option = SelectOption(label=role.name, value=str(role.id)) - options.append(option) + cache = self.cache.get(ctx.guild.id) + if cache: + role_id = cache.get(role) + else: + await ctx.send("Something went wrong, please try a different role", ephemeral=True) + return - select = Select( - options=options, - custom_id="to_delete", - placeholder="Select roles to remove", - min_values=1, - max_values=len(options), + setting.value.remove(role_id) + await setting.commit() + role = await ctx.guild.fetch_role(role_id) + if not role: + await ctx.send("Role not found in guild", ephemeral=True) + + remaining = [] + to_remove = [] + for id_ in setting.value: + if role := await ctx.guild.fetch_role(id_): + remaining.append(role) + else: + to_remove.append(id_) + + setting.value = [x for x in setting.value if x not in to_remove] + await setting.commit() + + fields = [ + EmbedField(name="Removed Role", value=role.mention), + EmbedField(name="Remaining Role(s)", value="\n".join([x.mention for x in remaining])), + ] + + embed = build_embed( + title="Rolegiver Updated", description="Role removed from rolegiver", fields=fields ) - components = [ActionRow(select)] - message = await ctx.send(content="\u200b", components=components) - try: - context = await self.bot.wait_for_component( - check=lambda x: ctx.author.id == x.context.author.id, - messages=message, - timeout=60 * 1, - ) - removed_roles = [] - for to_delete in context.context.values: - role = await ctx.guild.fetch_role(to_delete) - if role: - removed_roles.append(role) - setting.roles.remove(int(to_delete)) - await setting.commit() + embed.set_thumbnail(url=ctx.guild.icon.url) - for row in components: - for component in row.components: - component.disabled = True + embed.set_footer(text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}") - roles = [] - for role_id in setting.roles: - e_role = await ctx.guild.fetch_role(role_id) - if not e_role: - continue - roles.append(e_role) + await ctx.send( + embeds=embed, + ) - if roles: - roles.sort(key=lambda x: -x.position) - - value = "\n".join([r.mention for r in roles]) if 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="Rolegiver Updated", - description="Role removed from rolegiver", - fields=fields, - ) - - embed.set_thumbnail(url=ctx.guild.icon.url) - - embed.set_footer( - text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}" - ) - - await context.context.edit_origin( - content=f"Removed {len(context.context.values)} role(s)", - embeds=embed, - components=components, - ) - except asyncio.TimeoutError: - for row in components: - for component in row.components: - component.disabled = True - await message.edit(components=components) + if ctx.guild.id in self.cache: + self.cache[ctx.guild.id].pop(role.name) @rolegiver.subcommand(sub_cmd_name="list", sub_cmd_description="List rolegiver roles") async def _rolegiver_list(self, ctx: InteractionContext) -> None: @@ -382,6 +386,23 @@ class RolegiverCog(Extension): await setting.commit() await ctx.send("Rolegiver cleanup finished") + self.cache.pop(ctx.guild.id, None) + + @_rolegiver_remove.autocomplete("role") + async def _autocomplete(self, ctx: AutocompleteContext, role: str) -> None: + if not self.cache.get(ctx.guild.id): + rolegivers = await Rolegiver.find(q(guild=ctx.guild.id)).to_list(None) + for rolegiver in rolegivers: + role = await ctx.guild.fetch_role(rolegiver.role) + if not role: + await rolegiver.delete() + continue + if ctx.guild.id not in self.cache: + self.cache[ctx.guild.id] = {} + self.cache[ctx.guild.id][role.name] = role.id + results = process.extract(role, self.cache.get(ctx.guild.id).keys(), limit=25) + choices = [{"name": r[0], "value": r[0]} for r in results] + await ctx.send(choices=choices) def setup(bot: Client) -> None: diff --git a/jarvis/cogs/util.py b/jarvis/cogs/util.py index 37adfa1..b00dce8 100644 --- a/jarvis/cogs/util.py +++ b/jarvis/cogs/util.py @@ -52,6 +52,13 @@ class UtilCog(Extension): bot = SlashCommand(name="bot", description="Bot commands") + @bot.subcommand(sub_cmd_name="sex", sub_cmd_description="Have sex with JARVIS") + async def _sex(self, ctx: InteractionContext) -> None: + if ctx.author.id == 264072583987593217: + await ctx.send("Oh fuck no, go fuck yourself") + else: + await ctx.send("Not at this time, thank you for offering") + @bot.subcommand(sub_cmd_name="status", sub_cmd_description="Retrieve JARVIS status") @cooldown(bucket=Buckets.CHANNEL, rate=1, interval=30) async def _status(self, ctx: InteractionContext) -> None: