From 8fb9045254274ebaf32abe7ab08721f83dfb73a4 Mon Sep 17 00:00:00 2001 From: Shivansh-007 Date: Thu, 25 Feb 2021 11:18:33 +0530 Subject: [PATCH 1/2] Add support for top level root aliase - reference taken from Mark - python-discord/bot/pull/1124 --- bot/__init__.py | 9 ++++++++ bot/bot.py | 43 +++++++++++++++++++++++++++++++++++++++ bot/command.py | 23 +++++++++++++++++++++ bot/exts/github/github.py | 15 +++++++++----- 4 files changed, 85 insertions(+), 5 deletions(-) create mode 100644 bot/command.py diff --git a/bot/__init__.py b/bot/__init__.py index 24a5bfb1..4fc92b01 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -2,8 +2,11 @@ import os import typing +from functools import partial, partialmethod import loguru +from bot.command import Command +from discord.ext import commands from loguru import logger from .constants import LOG_FILE @@ -24,3 +27,9 @@ def should_rotate(message: loguru.Message, file: typing.TextIO) -> bool: logger.add(LOG_FILE, rotation=should_rotate) logger.info("Logging Process Started") + + +# 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) diff --git a/bot/bot.py b/bot/bot.py index 9c51a2c7..9845b410 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -1,4 +1,5 @@ import os +from typing import Optional from aiohttp import ClientSession from discord import Embed @@ -55,3 +56,45 @@ async def close(self) -> None: if self.http_session: await self.http_session.close() + + 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 + + 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: + # If alias matches any of the commands in the bot, then raise CommandRegistrationError + 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) diff --git a/bot/command.py b/bot/command.py new file mode 100644 index 00000000..72766561 --- /dev/null +++ b/bot/command.py @@ -0,0 +1,23 @@ +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. + + Example: + `!gh issue 72` can also be called as `!issue 72`. + `!gh src eval` can also be called as `!src eval`. + """ + + 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 must bbe of type list or tuple, containing strings." + ) diff --git a/bot/exts/github/github.py b/bot/exts/github/github.py index 143c9007..37f433b8 100644 --- a/bot/exts/github/github.py +++ b/bot/exts/github/github.py @@ -1,5 +1,6 @@ import typing +from bot.bot import Bot from bot.constants import BOT_REPO_URL from discord import Embed from discord.ext import commands @@ -18,7 +19,7 @@ class Github(commands.Cog): └ source Displays information about the bot's source code. """ - def __init__(self, bot: commands.Bot) -> None: + def __init__(self, bot: Bot) -> None: self.bot = bot @commands.group(name="github", aliases=("gh",), invoke_without_command=True) @@ -26,7 +27,7 @@ async def github_group(self, ctx: commands.Context) -> None: """Commands for Github.""" await ctx.send_help(ctx.command) - @github_group.command(name="profile") + @github_group.command(name="profile", root_aliases=("profile",)) @commands.cooldown(1, 10, BucketType.user) async def profile(self, ctx: commands.Context, username: str) -> None: """ @@ -39,7 +40,7 @@ async def profile(self, ctx: commands.Context, username: str) -> None: await ctx.send(embed=embed) - @github_group.command(name="issue", aliases=("pr",)) + @github_group.command(name="issue", aliases=("pr",), root_aliases=("pr", "issue")) async def issue( self, ctx: commands.Context, @@ -63,7 +64,11 @@ async def issue( await ctx.send(embed=embed) - @github_group.command(name="source", aliases=("src", "inspect")) + @github_group.command( + name="source", + aliases=("src", "inspect"), + root_aliases=("src", "inspect", "source"), + ) async def source_command( self, ctx: commands.Context, *, source_item: typing.Optional[str] = None ) -> None: @@ -86,6 +91,6 @@ async def source_command( await ctx.send(embed=embed) -def setup(bot: commands.Bot) -> None: +def setup(bot: Bot) -> None: """Load the Github cog.""" bot.add_cog(Github(bot)) From dddb68b2f56f9b234b63531532b9684dea72f303 Mon Sep 17 00:00:00 2001 From: Shivansh-007 Date: Thu, 25 Feb 2021 11:41:04 +0530 Subject: [PATCH 2/2] Similary add top-level groups for a groups's group. Well I know this sounds confusing, so go have a look at bot.group.Group class's docstring and it would be pretty clear to you --- bot/__init__.py | 4 ++++ bot/command.py | 2 +- bot/group.py | 23 +++++++++++++++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) create mode 100644 bot/group.py diff --git a/bot/__init__.py b/bot/__init__.py index 4fc92b01..a71b05f4 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -6,6 +6,7 @@ import loguru from bot.command import Command +from bot.group import Group from discord.ext import commands from loguru import logger @@ -33,3 +34,6 @@ def should_rotate(message: loguru.Message, file: typing.TextIO) -> bool: # Must be patched before any cogs are added. commands.command = partial(commands.command, cls=Command) commands.GroupMixin.command = partialmethod(commands.GroupMixin.command, cls=Command) + +commands.group = partial(commands.group, cls=Group) +commands.GroupMixin.group = partialmethod(commands.GroupMixin.group, cls=Group) diff --git a/bot/command.py b/bot/command.py index 72766561..b6dc6e7c 100644 --- a/bot/command.py +++ b/bot/command.py @@ -19,5 +19,5 @@ def __init__(self, *args, **kwargs): if not isinstance(self.root_aliases, (list, tuple)): raise TypeError( - "Root aliases must bbe of type list or tuple, containing strings." + "Root aliases must be of type list or tuple, containing strings." ) diff --git a/bot/group.py b/bot/group.py new file mode 100644 index 00000000..6bd18966 --- /dev/null +++ b/bot/group.py @@ -0,0 +1,23 @@ +from discord.ext import commands + + +class Group(commands.Group): + """ + A `discord.ext.commands.Group` 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 groups rather than being aliases of the group's group + + Example: + `!cmd gh issue 72` can also be called as `!gh issue 72`. + `!cmd src eval` can also be called as `!gh src eval`. + """ + + 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 must be of type list or tuple, containing strings." + )