Add root alias support for commands

1 files changed, 41 insertions, 0 deletions

diff --git a/bot/bot.py b/bot/bot.py
index 97b09243..f51139bb 100644
--- a/bot/bot.py
+++ b/bot/bot.py
@@ -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]