diff options
Diffstat (limited to 'bot')
| -rw-r--r-- | bot/__init__.py | 9 | ||||
| -rw-r--r-- | bot/bot.py | 41 | ||||
| -rw-r--r-- | bot/command.py | 18 | ||||
| -rw-r--r-- | bot/exts/evergreen/help.py | 4 | ||||
| -rw-r--r-- | bot/exts/evergreen/pfp_modify.py | 3 |
5 files changed, 73 insertions, 2 deletions
diff --git a/bot/__init__.py b/bot/__init__.py index bdb18666..c8550537 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -2,10 +2,13 @@ import asyncio import logging import logging.handlers import os +from functools import partial, partialmethod from pathlib import Path import arrow +from discord.ext import commands +from bot.command import Command from bot.constants import Client @@ -70,3 +73,9 @@ logging.getLogger().info('Logging initialization complete') # On Windows, the selector event loop is required for aiodns. if os.name == "nt": asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) + + +# Monkey-patch discord.py decorators to use the Command subclass which supports root aliases. +# Must be patched before any cogs are added. +commands.command = partial(commands.command, cls=Command) +commands.GroupMixin.command = partialmethod(commands.GroupMixin.command, cls=Command) @@ -64,6 +64,26 @@ class Bot(commands.Bot): super().add_cog(cog) log.info(f"Cog loaded: {cog.qualified_name}") + def add_command(self, command: commands.Command) -> None: + """Add `command` as normal and then add its root aliases to the bot.""" + super().add_command(command) + self._add_root_aliases(command) + + def remove_command(self, name: str) -> Optional[commands.Command]: + """ + Remove a command/alias as normal and then remove its root aliases from the bot. + + Individual root aliases cannot be removed by this function. + To remove them, either remove the entire command or manually edit `bot.all_commands`. + """ + command = super().remove_command(name) + if command is None: + # Even if it's a root alias, there's no way to get the Bot instance to remove the alias. + return + + self._remove_root_aliases(command) + return command + async def on_command_error(self, context: commands.Context, exception: DiscordException) -> None: """Check command errors for UserInputError and reset the cooldown if thrown.""" if isinstance(exception, commands.UserInputError): @@ -124,6 +144,27 @@ class Bot(commands.Bot): """ await self._guild_available.wait() + def _add_root_aliases(self, command: commands.Command) -> None: + """Recursively add root aliases for `command` and any of its subcommands.""" + if isinstance(command, commands.Group): + for subcommand in command.commands: + self._add_root_aliases(subcommand) + + for alias in getattr(command, "root_aliases", ()): + if alias in self.all_commands: + raise commands.CommandRegistrationError(alias, alias_conflict=True) + + self.all_commands[alias] = command + + def _remove_root_aliases(self, command: commands.Command) -> None: + """Recursively remove root aliases for `command` and any of its subcommands.""" + if isinstance(command, commands.Group): + for subcommand in command.commands: + self._remove_root_aliases(subcommand) + + for alias in getattr(command, "root_aliases", ()): + self.all_commands.pop(alias, None) + _allowed_roles = [discord.Object(id_) for id_ in constants.MODERATION_ROLES] diff --git a/bot/command.py b/bot/command.py new file mode 100644 index 00000000..0fb900f7 --- /dev/null +++ b/bot/command.py @@ -0,0 +1,18 @@ +from discord.ext import commands + + +class Command(commands.Command): + """ + A `discord.ext.commands.Command` subclass which supports root aliases. + + A `root_aliases` keyword argument is added, which is a sequence of alias names that will act as + top-level commands rather than being aliases of the command's group. It's stored as an attribute + also named `root_aliases`. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.root_aliases = kwargs.get("root_aliases", []) + + if not isinstance(self.root_aliases, (list, tuple)): + raise TypeError("Root aliases of a command must be a list or a tuple of strings.") diff --git a/bot/exts/evergreen/help.py b/bot/exts/evergreen/help.py index 91147243..f557e42e 100644 --- a/bot/exts/evergreen/help.py +++ b/bot/exts/evergreen/help.py @@ -289,7 +289,9 @@ class HelpSession: parent = self.query.full_parent_name + ' ' if self.query.parent else '' paginator.add_line(f'**```{prefix}{parent}{signature}```**') - aliases = ', '.join(f'`{a}`' for a in self.query.aliases) + aliases = [f"`{alias}`" if not parent else f"`{parent} {alias}`" for alias in self.query.aliases] + aliases += [f"`{alias}`" for alias in getattr(self.query, "root_aliases", ())] + aliases = ", ".join(sorted(aliases)) if aliases: paginator.add_line(f'**Can also use:** {aliases}\n') diff --git a/bot/exts/evergreen/pfp_modify.py b/bot/exts/evergreen/pfp_modify.py index 4245432e..a3f7e3f8 100644 --- a/bot/exts/evergreen/pfp_modify.py +++ b/bot/exts/evergreen/pfp_modify.py @@ -139,12 +139,13 @@ class PfpModify(commands.Cog): await ctx.send(file=file, embed=embed) @commands.max_concurrency(1, commands.BucketType.guild, wait=True) + @commands.group() async def pfp_modify(self, ctx: commands.Context) -> None: """Groups all of the pfp modifing commands to allow a single concurrency limit.""" if not ctx.invoked_subcommand: await ctx.send_help(ctx.command) - @pfp_modify.command(name="8bitify") + @pfp_modify.command(name="8bitify", root_aliases=("8bitify",)) async def eightbit_command(self, ctx: commands.Context) -> None: """Pixelates your avatar and changes the palette to an 8bit one.""" async with ctx.typing(): |