Support app command contexts

discord/app_commands/commands.py

@@ -48,7 +48,9 @@
 import re
 from copy import copy as shallow_copy
 
+
 from ..enums import AppCommandOptionType, AppCommandType, ChannelType, Locale
+from ..flags import AppCommandContext
 from .models import Choice
 from .transformers import annotation_to_parameter, CommandParameter, NoneType
 from .errors import AppCommandError, CheckFailure, CommandInvokeError, CommandSignatureMismatch, CommandAlreadyRegistered
@@ -87,6 +89,9 @@
     'autocomplete',
     'guilds',
     'guild_only',
+    'dm_only',
+    'private_channel_only',
+    'allow_contexts',
     'default_permissions',
 )
 
@@ -618,6 +623,9 @@ class Command(Generic[GroupT, P, T]):
         Whether the command should only be usable in guild contexts.
 
         Due to a Discord limitation, this does not work on subcommands.
+    allowed_contexts: Optional[:class:`~discord.flags.AppCommandContext`]
+        The contexts that the command is allowed to be used in.
+        Overrides ``guild_only`` if this is set.
     nsfw: :class:`bool`
         Whether the command is NSFW and should only work in NSFW channels.
 
@@ -638,6 +646,7 @@ def __init__(
         nsfw: bool = False,
         parent: Optional[Group] = None,
         guild_ids: Optional[List[int]] = None,
+        allowed_contexts: Optional[AppCommandContext] = None,
         auto_locale_strings: bool = True,
         extras: Dict[Any, Any] = MISSING,
     ):
@@ -672,6 +681,9 @@ def __init__(
             callback, '__discord_app_commands_default_permissions__', None
         )
         self.guild_only: bool = getattr(callback, '__discord_app_commands_guild_only__', False)
+        self.allowed_contexts: Optional[AppCommandContext] = allowed_contexts or getattr(
+            callback, '__discord_app_commands_contexts__', None
+        )
         self.nsfw: bool = nsfw
         self.extras: Dict[Any, Any] = extras or {}
 
@@ -760,6 +772,7 @@ def to_dict(self) -> Dict[str, Any]:
             base['nsfw'] = self.nsfw
             base['dm_permission'] = not self.guild_only
             base['default_member_permissions'] = None if self.default_permissions is None else self.default_permissions.value
+            base['contexts'] = self.allowed_contexts.to_array() if self.allowed_contexts is not None else None
 
         return base
 
@@ -1167,6 +1180,9 @@ class ContextMenu:
     guild_only: :class:`bool`
         Whether the command should only be usable in guild contexts.
         Defaults to ``False``.
+    allowed_contexts: Optional[:class:`~discord.flags.AppCommandContext`]
+        The contexts that this context menu is allowed to be used in.
+        Overrides ``guild_only`` if set.
     nsfw: :class:`bool`
         Whether the command is NSFW and should only work in NSFW channels.
         Defaults to ``False``.
@@ -1189,6 +1205,7 @@ def __init__(
         type: AppCommandType = MISSING,
         nsfw: bool = False,
         guild_ids: Optional[List[int]] = None,
+        allowed_contexts: Optional[AppCommandContext] = MISSING,
         auto_locale_strings: bool = True,
         extras: Dict[Any, Any] = MISSING,
     ):
@@ -1214,6 +1231,9 @@ def __init__(
         )
         self.nsfw: bool = nsfw
         self.guild_only: bool = getattr(callback, '__discord_app_commands_guild_only__', False)
+        self.allowed_contexts: Optional[AppCommandContext] = allowed_contexts or getattr(
+            callback, '__discord_app_commands_contexts__', None
+        )
         self.checks: List[Check] = getattr(callback, '__discord_app_commands_checks__', [])
         self.extras: Dict[Any, Any] = extras or {}
 
@@ -1249,6 +1269,7 @@ def to_dict(self) -> Dict[str, Any]:
             'name': self.name,
             'type': self.type.value,
             'dm_permission': not self.guild_only,
+            'contexts': self.allowed_contexts.to_array() if self.allowed_contexts is not None else None,
             'default_member_permissions': None if self.default_permissions is None else self.default_permissions.value,
             'nsfw': self.nsfw,
         }
@@ -1405,6 +1426,9 @@ class shortened to 100 characters.
         Whether the group should only be usable in guild contexts.
 
         Due to a Discord limitation, this does not work on subcommands.
+    allowed_contexts: Optional[:class:`~discord.flags.AppCommandContext`]
+        The contexts that this group is allowed to be used in. Overrides
+        guild_only if set.
     nsfw: :class:`bool`
         Whether the command is NSFW and should only work in NSFW channels.
 
@@ -1424,6 +1448,7 @@ class shortened to 100 characters.
     __discord_app_commands_group_locale_description__: Optional[locale_str] = None
     __discord_app_commands_group_nsfw__: bool = False
     __discord_app_commands_guild_only__: bool = MISSING
+    __discord_app_commands_contexts__: Optional[AppCommandContext] = MISSING
     __discord_app_commands_default_permissions__: Optional[Permissions] = MISSING
     __discord_app_commands_has_module__: bool = False
     __discord_app_commands_error_handler__: Optional[
@@ -1492,6 +1517,7 @@ def __init__(
         parent: Optional[Group] = None,
         guild_ids: Optional[List[int]] = None,
         guild_only: bool = MISSING,
+        allowed_contexts: Optional[AppCommandContext] = MISSING,
         nsfw: bool = MISSING,
         auto_locale_strings: bool = True,
         default_permissions: Optional[Permissions] = MISSING,
@@ -1540,6 +1566,14 @@ def __init__(
 
         self.guild_only: bool = guild_only
 
+        if allowed_contexts is MISSING:
+            if cls.__discord_app_commands_contexts__ is MISSING:
+                allowed_contexts = None
+            else:
+                allowed_contexts = cls.__discord_app_commands_contexts__
+
+        self.allowed_contexts: Optional[AppCommandContext] = allowed_contexts
+
         if nsfw is MISSING:
             nsfw = cls.__discord_app_commands_group_nsfw__
 
@@ -1669,6 +1703,7 @@ def to_dict(self) -> Dict[str, Any]:
             base['nsfw'] = self.nsfw
             base['dm_permission'] = not self.guild_only
             base['default_member_permissions'] = None if self.default_permissions is None else self.default_permissions.value
+            base['contexts'] = self.allowed_contexts.to_array() if self.allowed_contexts is not None else None
 
         return base
 
@@ -2418,8 +2453,16 @@ async def my_guild_only_command(interaction: discord.Interaction) -> None:
     def inner(f: T) -> T:
         if isinstance(f, (Command, Group, ContextMenu)):
             f.guild_only = True
+            allowed_contexts = f.allowed_contexts or AppCommandContext.none()
+            f.allowed_contexts = allowed_contexts
         else:
             f.__discord_app_commands_guild_only__ = True  # type: ignore # Runtime attribute assignment
+
+            allowed_contexts = getattr(f, '__discord_app_commands_contexts__', None) or AppCommandContext.none()
+            f.__discord_app_commands_contexts__ = allowed_contexts  # type: ignore # Runtime attribute assignment
+
+        allowed_contexts.guild = True
+
         return f
 
     # Check if called with parentheses or not
@@ -2430,6 +2473,135 @@ def inner(f: T) -> T:
         return inner(func)
 
 
+def private_channel_only(func: Optional[T] = None) -> Union[T, Callable[[T], T]]:
+    """A decorator that indicates this command can only be used in the context of DMs and group DMs.
+
+    This is **not** implemented as a :func:`check`, and is instead verified by Discord server side.
+    Therefore, there is no error handler called when a command is used within a guild.
+
+    This decorator can be called with or without parentheses.
+
+    Due to a Discord limitation, this decorator does nothing in subcommands and is ignored.
+
+    Examples
+    ---------
+
+    .. code-block:: python3
+
+        @app_commands.command()
+        @app_commands.private_channel_only()
+        async def my_private_channel_only_command(interaction: discord.Interaction) -> None:
+            await interaction.response.send_message('I am only available in DMs and GDMs!')
+    """
+
+    def inner(f: T) -> T:
+        if isinstance(f, (Command, Group, ContextMenu)):
+            f.guild_only = False
+            allowed_contexts = f.allowed_contexts or AppCommandContext.none()
+            f.allowed_contexts = allowed_contexts
+        else:
+            allowed_contexts = getattr(f, '__discord_app_commands_contexts__', None) or AppCommandContext.none()
+            f.__discord_app_commands_contexts__ = allowed_contexts  # type: ignore # Runtime attribute assignment
+
+        allowed_contexts.private_channel = True
+
+        return f
+
+    # Check if called with parentheses or not
+    if func is None:
+        # Called with parentheses
+        return inner
+    else:
+        return inner(func)
+
+
+def dm_only(func: Optional[T] = None) -> Union[T, Callable[[T], T]]:
+    """A decorator that indicates this command can only be used in the context of bot DMs.
+
+    This is **not** implemented as a :func:`check`, and is instead verified by Discord server side.
+    Therefore, there is no error handler called when a command is used within a guild or group DM.
+
+    This decorator can be called with or without parentheses.
+
+    Due to a Discord limitation, this decorator does nothing in subcommands and is ignored.
+
+    Examples
+    ---------
+
+    .. code-block:: python3
+
+        @app_commands.command()
+        @app_commands.dm_only()
+        async def my_dm_only_command(interaction: discord.Interaction) -> None:
+            await interaction.response.send_message('I am only available in DMs!')
+    """
+
+    def inner(f: T) -> T:
+        if isinstance(f, (Command, Group, ContextMenu)):
+            f.guild_only = False
+            allowed_contexts = f.allowed_contexts or AppCommandContext.none()
+            f.allowed_contexts = allowed_contexts
+        else:
+            allowed_contexts = getattr(f, '__discord_app_commands_contexts__', None) or AppCommandContext.none()
+            f.__discord_app_commands_contexts__ = allowed_contexts  # type: ignore # Runtime attribute assignment
+
+        allowed_contexts.dm_channel = True
+
+        return f
+
+    # Check if called with parentheses or not
+    if func is None:
+        # Called with parentheses
+        return inner
+    else:
+        return inner(func)
+
+
+# wrapper over previous 3 commands
+def allow_contexts(
+    guilds: bool = MISSING, dms: bool = MISSING, private_channels: bool = MISSING
+) -> Union[T, Callable[[T], T]]:
+    """A decorator that indicates this command can only be used in certain contexts.
+    Valid contexts are guilds, DMs and private channels.
+
+    This is **not** implemented as a :func:`check`, and is instead verified by Discord server side.
+
+    Due to a Discord limitation, this decorator does nothing in subcommands and is ignored.
+
+    Examples
+    ---------
+
+    .. code-block:: python3
+
+        @app_commands.command()
+        @app_commands.allow_contexts(guilds=False, dms=False, private_channels=True)
+        async def my_command(interaction: discord.Interaction) -> None:
+            await interaction.response.send_message('I am only available in guilds and private channels!')
+    """
+
+    def inner(f: T) -> T:
+        if isinstance(f, (Command, Group, ContextMenu)):
+            f.guild_only = False
+            allowed_contexts = f.allowed_contexts or AppCommandContext.none()
+            f.allowed_contexts = allowed_contexts
+        else:
+            allowed_contexts = getattr(f, '__discord_app_commands_contexts__', None) or AppCommandContext.none()
+            f.__discord_app_commands_contexts__ = allowed_contexts  # type: ignore # Runtime attribute assignment
+
+        if guilds is not MISSING:
+            allowed_contexts.guild = guilds
+
+        if dms is not MISSING:
+            allowed_contexts.dm_channel = dms
+
+        if private_channels is not MISSING:
+            allowed_contexts.private_channel = private_channels
+
+        return f
+
+    return inner
+
+
 def default_permissions(**perms: bool) -> Callable[[T], T]:
     r"""A decorator that sets the default permissions needed to execute this command.
 

discord/app_commands/models.py

@@ -26,9 +26,17 @@
 from datetime import datetime
 
 from .errors import MissingApplicationID
+from ..flags import AppCommandContext
 from .translator import TranslationContextLocation, TranslationContext, locale_str, Translator
 from ..permissions import Permissions
-from ..enums import AppCommandOptionType, AppCommandType, AppCommandPermissionType, ChannelType, Locale, try_enum
+from ..enums import (
+    AppCommandOptionType,
+    AppCommandType,
+    AppCommandPermissionType,
+    ChannelType,
+    Locale,
+    try_enum,
+)
 from ..mixins import Hashable
 from ..utils import _get_as_snowflake, parse_time, snowflake_time, MISSING
 from ..object import Object
@@ -160,6 +168,8 @@ class AppCommand(Hashable):
         The default member permissions that can run this command.
     dm_permission: :class:`bool`
         A boolean that indicates whether this command can be run in direct messages.
+    allowed_contexts: Optional[:class:`~discord.flags.AppCommandContext`]
+        The contexts that this command is allowed to be used in. Overrides the ``dm_permission`` attribute.
     guild_id: Optional[:class:`int`]
         The ID of the guild this command is registered in. A value of ``None``
         denotes that it is a global command.
@@ -179,6 +189,7 @@ class AppCommand(Hashable):
         'options',
         'default_member_permissions',
         'dm_permission',
+        'allowed_contexts',
         'nsfw',
         '_state',
     )
@@ -210,6 +221,13 @@ def _from_data(self, data: ApplicationCommandPayload) -> None:
             dm_permission = True
 
         self.dm_permission: bool = dm_permission
+
+        allowed_contexts = data.get('contexts')
+        if allowed_contexts is None:
+            self.allowed_contexts: Optional[AppCommandContext] = None
+        else:
+            self.allowed_contexts = AppCommandContext._from_value(allowed_contexts)
+
         self.nsfw: bool = data.get('nsfw', False)
         self.name_localizations: Dict[Locale, str] = _to_locale_dict(data.get('name_localizations') or {})
         self.description_localizations: Dict[Locale, str] = _to_locale_dict(data.get('description_localizations') or {})
@@ -223,6 +241,7 @@ def to_dict(self) -> ApplicationCommandPayload:
             'description': self.description,
             'name_localizations': {str(k): v for k, v in self.name_localizations.items()},
             'description_localizations': {str(k): v for k, v in self.description_localizations.items()},
+            'contexts': self.allowed_contexts.to_array() if self.allowed_contexts is not None else None,
             'options': [opt.to_dict() for opt in self.options],
         }  # type: ignore # Type checker does not understand this literal.