From 25b4bc277bd8ab18a62feb5f69096114f9fa760b Mon Sep 17 00:00:00 2001
From: Rapptz <rapptz@gmail.com>
Date: Tue, 8 Mar 2022 00:48:24 -0500
Subject: [PATCH] Add app_commands.guilds to set the guilds of a command in
 another way

This is mostly preparation for interopability with commands.Cog as this
would allow authors to specify the guilds for their cog defined
commands.
---
 discord/app_commands/commands.py | 48 ++++++++++++++++++++++++++++++++
 discord/app_commands/tree.py     | 19 +++++++++----
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/discord/app_commands/commands.py b/discord/app_commands/commands.py
index 5b8de52de..c3f06c4d9 100644
--- a/discord/app_commands/commands.py
+++ b/discord/app_commands/commands.py
@@ -56,6 +56,7 @@ from ..utils import resolve_annotation, MISSING, is_inside_class
 
 if TYPE_CHECKING:
     from typing_extensions import ParamSpec, Concatenate
+    from ..abc import Snowflake
     from .namespace import Namespace
     from .models import ChoiceT
 
@@ -68,6 +69,7 @@ __all__ = (
     'describe',
     'choices',
     'autocomplete',
+    'guilds',
 )
 
 if TYPE_CHECKING:
@@ -1067,3 +1069,49 @@ def autocomplete(**parameters: AutocompleteCallback[GroupT, ChoiceT]) -> Callabl
         return inner
 
     return decorator
+
+
+def guilds(*guild_ids: Union[Snowflake, int]) -> Callable[[T], T]:
+    r"""Associates the given guilds with the command.
+
+    When the command instance is added to a :class:`CommandTree`, the guilds that are
+    specified by this decorator become the default guilds that it's added to rather
+    than being a global command.
+
+    .. note::
+
+        Due to an implementation quirk and Python limitation, if this is used in conjunction
+        with the :meth:`CommandTree.command` or :meth:`CommandTree.context_menu` decorator
+        then this must go below that decorator.
+
+    Example:
+
+    .. code-block:: python3
+
+            MY_GUILD_ID = discord.Object(...)  # Guild ID here
+
+            @app_commands.command()
+            @app_commands.guilds(MY_GUILD_ID)
+            async def bonk(interaction: discord.Interaction):
+                await interaction.response.send_message('Bonk', ephemeral=True)
+
+    Parameters
+    -----------
+    \*guild_ids: Union[:class:`int`, :class:`~discord.abc.Snowflake`]
+        The guilds to associate this command with. The command tree will
+        use this as the default when added rather than adding it as a global
+        command.
+    """
+
+    defaults: List[int] = [g if isinstance(g, int) else g.id for g in guild_ids]
+
+    def decorator(inner: T) -> T:
+        if isinstance(inner, Command):
+            inner._callback.__discord_app_commands_default_guilds__ = defaults
+        else:
+            # Runtime attribute assignment
+            inner.__discord_app_commands_default_guilds__ = defaults  # type: ignore
+
+        return inner
+
+    return decorator
diff --git a/discord/app_commands/tree.py b/discord/app_commands/tree.py
index dadcf3dad..24263a744 100644
--- a/discord/app_commands/tree.py
+++ b/discord/app_commands/tree.py
@@ -26,7 +26,7 @@ from __future__ import annotations
 import inspect
 import sys
 import traceback
-from typing import Callable, Dict, Generic, List, Literal, Optional, TYPE_CHECKING, Set, Tuple, TypeVar, Union, overload
+from typing import Any, Callable, Dict, Generic, List, Literal, Optional, TYPE_CHECKING, Set, Tuple, TypeVar, Union, overload
 
 
 from .namespace import Namespace, ResolveKey
@@ -54,14 +54,23 @@ __all__ = ('CommandTree',)
 ClientT = TypeVar('ClientT', bound='Client')
 
 
-def _retrieve_guild_ids(guild: Optional[Snowflake] = MISSING, guilds: List[Snowflake] = MISSING) -> Optional[Set[int]]:
+def _retrieve_guild_ids(
+    callback: Any, guild: Optional[Snowflake] = MISSING, guilds: List[Snowflake] = MISSING
+) -> Optional[Set[int]]:
     if guild is not MISSING and guilds is not MISSING:
         raise TypeError('cannot mix guild and guilds keyword arguments')
 
-    # guilds=[] or guilds=[...] or no args at all
+    # guilds=[] or guilds=[...]
     if guild is MISSING:
-        if not guilds:
+        # If no arguments are given then it should default to the ones
+        # given to the guilds(...) decorator or None for global.
+        if guild is MISSING:
+            return getattr(callback, '__discord_app_commands_default_guilds__', None)
+
+        # guilds=[] is the same as global
+        if len(guilds) == 0:
             return None
+
         return {g.id for g in guilds}
 
     # At this point it should be...
@@ -176,7 +185,7 @@ class CommandTree(Generic[ClientT]):
             This is currently 100 for slash commands and 5 for context menu commands.
         """
 
-        guild_ids = _retrieve_guild_ids(guild, guilds)
+        guild_ids = _retrieve_guild_ids(getattr(command, '_callback', None), guild, guilds)
         if isinstance(command, ContextMenu):
             type = command.type.value
             name = command.name