Skip to content

Instantly share code, notes, and snippets.

@LeoCx1000
Last active August 25, 2024 15:24
Show Gist options
  • Save LeoCx1000/021dc52981299b95ea7790416e4f5ca4 to your computer and use it in GitHub Desktop.
Save LeoCx1000/021dc52981299b95ea7790416e4f5ca4 to your computer and use it in GitHub Desktop.
Mentionable CommandTree implementation to allow mentioning slash commands in discord.py

discord.py MentionableTree implementation

My implementation of a CommandTree subclass which allows for finding a mention for a given app command.

How to use?

Just copy the entire third file into a new file python of your project, and import it.

Additions to the default CommandTree.

await find_mention_for(command, *, guild=None)

Retrieves the mention of an AppCommand given a specific command object or name.

Parameters:

  • name ( Union[ app_commands.Command, commands.HybridCommand, str] ) - The command to retrieve the mention for.
  • guild (Optional[abc.Snowflake]) - The scope (guild) from which to retrieve the commands from. If None is given or not passed, only the global scope will be searched, however the global scope will also be searched if a guild is passed.

Returns:

  • (Optional[str]) The command mention, if found.

async for command, mention in walk_mentions(*, guild=None, yield_unknown=False):

Gets all valid mentions for app commands in a specific guild.

This takes into consideration group commands, it will only return mentions for the command's children, and not the parent as parents aren't mentionable.

:Parameters:

  • guild (Optional[abc.Snowflake]) - The guild to get commands for. If not given, it will only return global commands.
  • yield_unknown: (bool) - If this is set to True, the yielded mention can be None, instead of it being ignored. Defaults to False.

Yields:

Changelog

  • 2023-08-14: Turn find_mention_for into a coroutine, and make it auto-fetch commands.
  • 2024-03-12: Make find_mention_for respect fallback_to_global.
  • 2024-03-24: Fix this issue.
  • 2024-05-30: Redo some of the logic, add a cache for mentions, and add a walk_command_mentions(guild=...) method which yields tuple[Command, str]
from .mentionable_tree import *
# Why does this exist? In case you want to `git clone [GIST_LINK] folder_name` and `from folder_name import MentionableTree`
# fmt: off
from logging import getLogger
from typing import Optional, List, Generator
import discord
from discord import app_commands
from discord.ext import commands
__all__ = ('MentionableTree',)
_log = getLogger(__name__)
class MentionableTree(app_commands.CommandTree):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.application_commands: dict[Optional[int], List[app_commands.AppCommand]] = {}
self.cache: dict[Optional[int], dict[app_commands.Command | commands.HybridCommand | str, str]] = {}
async def sync(self, *, guild: Optional[discord.abc.Snowflake] = None):
"""Method overwritten to store the commands."""
ret = await super().sync(guild=guild)
guild_id = guild.id if guild else None
self.application_commands[guild_id] = ret
self.cache.pop(guild_id, None)
return ret
async def fetch_commands(self, *, guild: Optional[discord.abc.Snowflake] = None):
"""Method overwritten to store the commands."""
ret = await super().fetch_commands(guild=guild)
guild_id = guild.id if guild else None
self.application_commands[guild_id] = ret
self.cache.pop(guild_id, None)
return ret
async def get_or_fetch_commands(self, *, guild: Optional[discord.abc.Snowflake] = None):
"""Method overwritten to store the commands."""
try:
return self.application_commands[guild.id if guild else None]
except KeyError:
return await self.fetch_commands(guild=guild)
async def find_mention_for(
self,
command: app_commands.Command | commands.HybridCommand | str,
*,
guild: Optional[discord.abc.Snowflake] = None,
) -> Optional[str]:
"""Retrieves the mention of an AppCommand given a specific command name, and optionally, a guild.
Parameters
----------
name: Union[:class:`app_commands.Command`, :class:`commands.HybridCommand`, str]
The command to retrieve the mention for.
guild: Optional[:class:`discord.abc.Snowflake`]
The scope (guild) from which to retrieve the commands from. If None is given or not passed,
only the global scope will be searched, however the global scope will also be searched if
a guild is passed.
Returns
-------
str | None
The command mention, if found.
"""
guild_id = guild.id if guild else None
try:
return self.cache[guild_id][command]
except KeyError:
pass
# If a guild is given, and fallback to global is set to True, then we must also
# check the global scope, as commands for both show in a guild.
check_global = self.fallback_to_global is True and guild is not None
if isinstance(command, str):
# Try and find a command by that name. discord.py does not return children from tree.get_command, but
# using walk_commands and utils.get is a simple way around that.
_command = discord.utils.get(self.walk_commands(guild=guild), qualified_name=command)
if check_global and not _command:
_command = discord.utils.get(self.walk_commands(), qualified_name=command)
else:
_command = command
if not _command:
return None
local_commands = await self.get_or_fetch_commands(guild=guild)
app_command_found = discord.utils.get(local_commands, name=(_command.root_parent or _command).name)
if check_global and not app_command_found:
global_commands = await self.get_or_fetch_commands(guild=None)
app_command_found = discord.utils.get(global_commands, name=(_command.root_parent or _command).name)
if not app_command_found:
return None
mention = f"</{_command.qualified_name}:{app_command_found.id}>"
self.cache.setdefault(guild_id, {})
self.cache[guild_id][command] = mention
return mention
def _walk_children(self, commands: list[app_commands.Group | app_commands.Command]) -> Generator[app_commands.Command, None, None]:
for command in commands:
if isinstance(command, app_commands.Group):
yield from self._walk_children(command.commands)
else:
yield command
async def walk_mentions(self, *, guild: Optional[discord.abc.Snowflake] = None):
"""Gets all valid mentions for app commands in a specific guild.
This takes into consideration group commands, it will only return mentions for
the command's children, and not the parent as parents aren't mentionable.
Parameters
----------
guild: Optional[discord.Guild]
The guild to get commands for. If not given, it will only return global commands.
Yields
------
Tuple[Union[:class:`app_commands.Command`, :class:`commands.HybridCommand`], str]
"""
for command in self._walk_children(self.get_commands(guild=guild, type=discord.AppCommandType.chat_input)):
mention = await self.find_mention_for(command, guild=guild)
if mention:
yield command, mention
if guild and self.fallback_to_global is True:
for command in self._walk_children(self.get_commands(guild=None, type=discord.AppCommandType.chat_input)):
mention = await self.find_mention_for(command, guild=guild)
if mention:
yield command, mention
else:
_log.warn("Could not find a mention for command %s in the API. Are you out of sync?", command)
@Sachaa-Thanasius
Copy link

Minor bug:

check_global = self.fallback_to_global is True and guild is not None

This should be an or condition and not an and one, based on the description in the docstring.

@LeoCx1000
Copy link
Author

LeoCx1000 commented Mar 24, 2024

You're right, I suppose this could cause issues if you're looking for a global command's mention in a guild, but you have fallback_to_global set to False. I hadn't thought of this edge case, thanks.

Oh! I see the other logic issue now, oops.

@Sachaa-Thanasius
Copy link

Sachaa-Thanasius commented Mar 25, 2024

There may be more to it, but I didn’t look closer beyond the first supposed mismatch in logic vs. docs. Glad it was helpful. Thanks for the great snippet.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment