Add application command cooldown decorators

Fix #7790

discord/app_commands/__init__.py

@@ -16,3 +16,4 @@ from .tree import *
 from .namespace import *
 from .transformers import *
 from . import checks as checks
+from .checks import Cooldown as Cooldown

discord/app_commands/checks.py

@@ -25,12 +25,19 @@ DEALINGS IN THE SOFTWARE.
 from __future__ import annotations
 
 from typing import (
+    Any,
+    Coroutine,
+    Dict,
+    Hashable,
     Union,
     Callable,
     TypeVar,
+    Optional,
     TYPE_CHECKING,
 )
 
+import time
+
 from .commands import check
 from .errors import (
     NoPrivateMessage,
@@ -38,25 +45,150 @@ from .errors import (
     MissingAnyRole,
     MissingPermissions,
     BotMissingPermissions,
+    CommandOnCooldown,
 )
 
 from ..user import User
 from ..permissions import Permissions
-from ..utils import get as utils_get
+from ..utils import get as utils_get, MISSING, maybe_coroutine
 
 T = TypeVar('T')
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
     from ..interactions import Interaction
 
+    CooldownFunction = Union[
+        Callable[[Interaction], Coroutine[Any, Any, T]],
+        Callable[[Interaction], T],
+    ]
+
 __all__ = (
     'has_role',
     'has_any_role',
     'has_permissions',
     'bot_has_permissions',
+    'cooldown',
+    'dynamic_cooldown',
 )
 
 
+class Cooldown:
+    """Represents a cooldown for a command.
+
+    .. versionadded:: 2.0
+
+    Attributes
+    -----------
+    rate: :class:`float`
+        The total number of tokens available per :attr:`per` seconds.
+    per: :class:`float`
+        The length of the cooldown period in seconds.
+    """
+
+    __slots__ = ('rate', 'per', '_window', '_tokens', '_last')
+
+    def __init__(self, rate: float, per: float) -> None:
+        self.rate: int = int(rate)
+        self.per: float = float(per)
+        self._window: float = 0.0
+        self._tokens: int = self.rate
+        self._last: float = 0.0
+
+    def get_tokens(self, current: Optional[float] = None) -> int:
+        """Returns the number of available tokens before rate limiting is applied.
+
+        Parameters
+        ------------
+        current: Optional[:class:`float`]
+            The time in seconds since Unix epoch to calculate tokens at.
+            If not supplied then :func:`time.time()` is used.
+
+        Returns
+        --------
+        :class:`int`
+            The number of tokens available before the cooldown is to be applied.
+        """
+        if not current:
+            current = time.time()
+
+        tokens = self._tokens
+
+        if current > self._window + self.per:
+            tokens = self.rate
+        return tokens
+
+    def get_retry_after(self, current: Optional[float] = None) -> float:
+        """Returns the time in seconds until the cooldown will be reset.
+
+        Parameters
+        -------------
+        current: Optional[:class:`float`]
+            The current time in seconds since Unix epoch.
+            If not supplied, then :func:`time.time()` is used.
+
+        Returns
+        -------
+        :class:`float`
+            The number of seconds to wait before this cooldown will be reset.
+        """
+        current = current or time.time()
+        tokens = self.get_tokens(current)
+
+        if tokens == 0:
+            return self.per - (current - self._window)
+
+        return 0.0
+
+    def update_rate_limit(self, current: Optional[float] = None) -> Optional[float]:
+        """Updates the cooldown rate limit.
+
+        Parameters
+        -------------
+        current: Optional[:class:`float`]
+            The time in seconds since Unix epoch to update the rate limit at.
+            If not supplied, then :func:`time.time()` is used.
+
+        Returns
+        -------
+        Optional[:class:`float`]
+            The retry-after time in seconds if rate limited.
+        """
+        current = current or time.time()
+        self._last = current
+
+        self._tokens = self.get_tokens(current)
+
+        # first token used means that we start a new rate limit window
+        if self._tokens == self.rate:
+            self._window = current
+
+        # check if we are rate limited
+        if self._tokens == 0:
+            return self.per - (current - self._window)
+
+        # we're not so decrement our tokens
+        self._tokens -= 1
+
+    def reset(self) -> None:
+        """Reset the cooldown to its initial state."""
+        self._tokens = self.rate
+        self._last = 0.0
+
+    def copy(self) -> Self:
+        """Creates a copy of this cooldown.
+
+        Returns
+        --------
+        :class:`Cooldown`
+            A new instance of this cooldown.
+        """
+        return Cooldown(self.rate, self.per)
+
+    def __repr__(self) -> str:
+        return f'<Cooldown rate: {self.rate} per: {self.per} window: {self._window} tokens: {self._tokens}>'
+
+
 def has_role(item: Union[int, str], /) -> Callable[[T], T]:
     """A :func:`~discord.app_commands.check` that is added that checks if the member invoking the
     command has the role specified via the name or ID specified.
@@ -236,3 +368,134 @@ def bot_has_permissions(**perms: bool) -> Callable[[T], T]:
         raise BotMissingPermissions(missing)
 
     return check(predicate)
+
+
+def _create_cooldown_decorator(
+    key: CooldownFunction[Hashable], factory: CooldownFunction[Optional[Cooldown]]
+) -> Callable[[T], T]:
+
+    mapping: Dict[Any, Cooldown] = {}
+
+    async def get_bucket(
+        interaction: Interaction,
+        *,
+        mapping: Dict[Any, Cooldown] = mapping,
+        key: CooldownFunction[Hashable] = key,
+        factory: CooldownFunction[Optional[Cooldown]] = factory,
+    ) -> Optional[Cooldown]:
+        current = interaction.created_at.timestamp()
+        dead_keys = [k for k, v in mapping.items() if current > v._last + v.per]
+        for k in dead_keys:
+            del mapping[k]
+
+        k = await maybe_coroutine(key, interaction)
+        if k not in mapping:
+            bucket: Optional[Cooldown] = await maybe_coroutine(factory, interaction)
+            if bucket is not None:
+                mapping[k] = bucket
+        else:
+            bucket = mapping[k]
+
+        return bucket
+
+    async def predicate(interaction: Interaction) -> bool:
+        bucket = await get_bucket(interaction)
+        if bucket is None:
+            return True
+
+        retry_after = bucket.update_rate_limit(interaction.created_at.timestamp())
+        if retry_after is None:
+            return True
+
+        raise CommandOnCooldown(bucket, retry_after)
+
+    return check(predicate)
+
+
+def cooldown(
+    rate: float,
+    per: float,
+    *,
+    key: Optional[CooldownFunction[Hashable]] = MISSING,
+) -> Callable[[T], T]:
+    """A decorator that adds a cooldown to a command.
+
+    A cooldown allows a command to only be used a specific amount
+    of times in a specific time frame. These cooldowns are based off
+    of the ``key`` function provided. If a ``key`` is not provided
+    then it defaults to a user-level cooldown. The ``key`` function
+    must take a single parameter, the :class:`discord.Interaction` and
+    return a value that is used as a key to the internal cooldown mapping.
+
+    The ``key`` function can optionally be a coroutine.
+
+    If a cooldown is triggered, then :exc:`~discord.app_commands.CommandOnCooldown` is
+    raised to the error handlers.
+
+    Parameters
+    ------------
+    rate: :class:`int`
+        The number of times a command can be used before triggering a cooldown.
+    per: :class:`float`
+        The amount of seconds to wait for a cooldown when it's been triggered.
+    key: Optional[Callable[[:class:`discord.Interaction`], :class:`collections.abc.Hashable`]]
+        A function that returns a key to the mapping denoting the type of cooldown.
+        Can optionally be a coroutine. If not given then defaults to a user-level
+        cooldown. If ``None`` is passed then it is interpreted as a "global" cooldown.
+    """
+
+    if key is MISSING:
+        key_func = lambda interaction: interaction.user.id
+    elif key is None:
+        key_func = lambda i: None
+    else:
+        key_func = key
+
+    factory = lambda interaction: Cooldown(rate, per)
+
+    return _create_cooldown_decorator(key_func, factory)
+
+
+def dynamic_cooldown(
+    factory: CooldownFunction[Optional[Cooldown]],
+    *,
+    key: Optional[CooldownFunction[Hashable]] = MISSING,
+) -> Callable[[T], T]:
+    """A decorator that adds a dynamic cooldown to a command.
+
+    A cooldown allows a command to only be used a specific amount
+    of times in a specific time frame. These cooldowns are based off
+    of the ``key`` function provided. If a ``key`` is not provided
+    then it defaults to a user-level cooldown. The ``key`` function
+    must take a single parameter, the :class:`discord.Interaction` and
+    return a value that is used as a key to the internal cooldown mapping.
+
+    If a ``factory`` function is given, it must be a function that
+    accepts a single parameter of type :class:`discord.Interaction` and must
+    return a :class:`~discord.app_commands.Cooldown` or ``None``.
+    If ``None`` is returned then that cooldown is effectively bypassed.
+
+    Both ``key`` and ``factory`` can optionally be coroutines.
+
+    If a cooldown is triggered, then :exc:`~discord.app_commands.CommandOnCooldown` is
+    raised to the error handlers.
+
+    Parameters
+    ------------
+    factory: Optional[Callable[[:class:`discord.Interaction`], Optional[:class:`~discord.app_commands.Cooldown`]]]
+        A function that takes an interaction and returns a cooldown that will apply to that interaction
+        or ``None`` if the interaction should not have a cooldown.
+    key: Optional[Callable[[:class:`discord.Interaction`], :class:`collections.abc.Hashable`]]
+        A function that returns a key to the mapping denoting the type of cooldown.
+        Can optionally be a coroutine. If not given then defaults to a user-level
+        cooldown. If ``None`` is passed then it is interpreted as a "global" cooldown.
+    """
+
+    if key is MISSING:
+        key_func = lambda interaction: interaction.user.id
+    elif key is None:
+        key_func = lambda i: None
+    else:
+        key_func = key
+
+    return _create_cooldown_decorator(key_func, factory)

discord/app_commands/errors.py

@@ -43,12 +43,14 @@ __all__ = (
     'MissingAnyRole',
     'MissingPermissions',
     'BotMissingPermissions',
+    'CommandOnCooldown',
 )
 
 if TYPE_CHECKING:
     from .commands import Command, Group, ContextMenu
     from .transformers import Transformer
     from ..types.snowflake import Snowflake, SnowflakeList
+    from .checks import Cooldown
 
 
 class AppCommandError(DiscordException):
@@ -262,6 +264,27 @@ class BotMissingPermissions(CheckFailure):
         super().__init__(message, *args)
 
 
+class CommandOnCooldown(CheckFailure):
+    """An exception raised when the command being invoked is on cooldown.
+
+    This inherits from :exc:`~discord.app_commands.CheckFailure`.
+
+    .. versionadded:: 2.0
+
+    Attributes
+    -----------
+    cooldown: :class:`~discord.app_commands.Cooldown`
+        The cooldown that was triggered.
+    retry_after: :class:`float`
+        The amount of seconds to wait before you can retry again.
+    """
+
+    def __init__(self, cooldown: Cooldown, retry_after: float) -> None:
+        self.cooldown: Cooldown = cooldown
+        self.retry_after: float = retry_after
+        super().__init__(f'You are on cooldown. Try again in {retry_after:.2f}s')
+
+
 class CommandAlreadyRegistered(AppCommandError):
     """An exception raised when a command is already registered.
 

discord/ext/commands/cooldowns.py

@@ -33,6 +33,7 @@ from collections import deque
 
 from ...abc import PrivateChannel
 from .errors import MaxConcurrencyReached
+from discord.app_commands import Cooldown as Cooldown
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -79,120 +80,6 @@ class BucketType(Enum):
         return self.get_key(msg)
 
 
-class Cooldown:
-    """Represents a cooldown for a command.
-
-    Attributes
-    -----------
-    rate: :class:`int`
-        The total number of tokens available per :attr:`per` seconds.
-    per: :class:`float`
-        The length of the cooldown period in seconds.
-    """
-
-    __slots__ = ('rate', 'per', '_window', '_tokens', '_last')
-
-    def __init__(self, rate: float, per: float) -> None:
-        self.rate: int = int(rate)
-        self.per: float = float(per)
-        self._window: float = 0.0
-        self._tokens: int = self.rate
-        self._last: float = 0.0
-
-    def get_tokens(self, current: Optional[float] = None) -> int:
-        """Returns the number of available tokens before rate limiting is applied.
-
-        Parameters
-        ------------
-        current: Optional[:class:`float`]
-            The time in seconds since Unix epoch to calculate tokens at.
-            If not supplied then :func:`time.time()` is used.
-
-        Returns
-        --------
-        :class:`int`
-            The number of tokens available before the cooldown is to be applied.
-        """
-        if not current:
-            current = time.time()
-
-        tokens = self._tokens
-
-        if current > self._window + self.per:
-            tokens = self.rate
-        return tokens
-
-    def get_retry_after(self, current: Optional[float] = None) -> float:
-        """Returns the time in seconds until the cooldown will be reset.
-
-        Parameters
-        -------------
-        current: Optional[:class:`float`]
-            The current time in seconds since Unix epoch.
-            If not supplied, then :func:`time.time()` is used.
-
-        Returns
-        -------
-        :class:`float`
-            The number of seconds to wait before this cooldown will be reset.
-        """
-        current = current or time.time()
-        tokens = self.get_tokens(current)
-
-        if tokens == 0:
-            return self.per - (current - self._window)
-
-        return 0.0
-
-    def update_rate_limit(self, current: Optional[float] = None) -> Optional[float]:
-        """Updates the cooldown rate limit.
-
-        Parameters
-        -------------
-        current: Optional[:class:`float`]
-            The time in seconds since Unix epoch to update the rate limit at.
-            If not supplied, then :func:`time.time()` is used.
-
-        Returns
-        -------
-        Optional[:class:`float`]
-            The retry-after time in seconds if rate limited.
-        """
-        current = current or time.time()
-        self._last = current
-
-        self._tokens = self.get_tokens(current)
-
-        # first token used means that we start a new rate limit window
-        if self._tokens == self.rate:
-            self._window = current
-
-        # check if we are rate limited
-        if self._tokens == 0:
-            return self.per - (current - self._window)
-
-        # we're not so decrement our tokens
-        self._tokens -= 1
-
-    def reset(self) -> None:
-        """Reset the cooldown to its initial state."""
-        self._tokens = self.rate
-        self._last = 0.0
-
-    def copy(self) -> Cooldown:
-        """Creates a copy of this cooldown.
-
-        Returns
-        --------
-        :class:`Cooldown`
-            A new instance of this cooldown.
-        """
-        return Cooldown(self.rate, self.per)
-
-    def __repr__(self) -> str:
-        return f'<Cooldown rate: {self.rate} per: {self.per} window: {self._window} tokens: {self._tokens}>'
-
-
 class CooldownMapping:
     def __init__(
         self,

discord/ext/commands/core.py

@@ -2294,8 +2294,8 @@ def dynamic_cooldown(
 
     This differs from :func:`.cooldown` in that it takes a function that
     accepts a single parameter of type :class:`.discord.Message` and must
-    return a :class:`.Cooldown` or ``None``. If ``None`` is returned then
+    return a :class:`~discord.app_commands.Cooldown` or ``None``.
-    that cooldown is effectively bypassed.
+    If ``None`` is returned then that cooldown is effectively bypassed.
 
     A cooldown allows a command to only be used a specific amount
     of times in a specific time frame. These cooldowns can be based
@@ -2312,7 +2312,7 @@ def dynamic_cooldown(
 
     Parameters
     ------------
-    cooldown: Callable[[:class:`.discord.Message`], Optional[:class:`.Cooldown`]]
+    cooldown: Callable[[:class:`.discord.Message`], Optional[:class:`~discord.app_commands.Cooldown`]]
         A function that takes a message and returns a cooldown that will
         apply to this invocation or ``None`` if the cooldown should be bypassed.
     type: :class:`.BucketType`

discord/ext/commands/errors.py

@@ -587,7 +587,7 @@ class CommandOnCooldown(CommandError):
 
     Attributes
     -----------
-    cooldown: :class:`.Cooldown`
+    cooldown: :class:`~discord.app_commands.Cooldown`
         A class with attributes ``rate`` and ``per`` similar to the
         :func:`.cooldown` decorator.
     type: :class:`BucketType`

docs/ext/commands/api.rst

@@ -330,14 +330,6 @@ Checks
 
 .. _ext_commands_api_context:
 
-Cooldown
----------
-
-.. attributetable:: discord.ext.commands.Cooldown
-
-.. autoclass:: discord.ext.commands.Cooldown
-    :members:
-
 Context
 --------
 

docs/interactions/api.rst

@@ -451,7 +451,7 @@ Group
     :members:
 
 Decorators
-+++++++++++
+~~~~~~~~~~~
 
 .. autofunction:: discord.app_commands.command
     :decorator:
@@ -475,7 +475,7 @@ Decorators
     :decorator:
 
 Checks
-+++++++
+~~~~~~~
 
 .. autofunction:: discord.app_commands.check
     :decorator:
@@ -492,6 +492,19 @@ Checks
 .. autofunction:: discord.app_commands.checks.bot_has_permissions
     :decorator:
 
+.. autofunction:: discord.app_commands.checks.cooldown
+    :decorator:
+
+.. autofunction:: discord.app_commands.checks.dynamic_cooldown
+    :decorator:
+
+Cooldown
+~~~~~~~~~
+
+.. attributetable:: discord.app_commands.Cooldown
+
+.. autoclass:: discord.app_commands.Cooldown
+    :members:
 
 
 Namespace
@@ -559,6 +572,9 @@ Exceptions
 .. autoexception:: discord.app_commands.BotMissingPermissions
     :members:
 
+.. autoexception:: discord.app_commands.CommandOnCooldown
+    :members:
+
 .. autoexception:: discord.app_commands.CommandAlreadyRegistered
     :members:
 
@@ -583,6 +599,7 @@ Exception Hierarchy
                 - :exc:`~discord.app_commands.MissingAnyRole`
                 - :exc:`~discord.app_commands.MissingPermissions`
                 - :exc:`~discord.app_commands.BotMissingPermissions`
+                - :exc:`~discord.app_commands.CommandOnCooldown`
             - :exc:`~discord.app_commands.CommandAlreadyRegistered`
             - :exc:`~discord.app_commands.CommandSignatureMismatch`
             - :exc:`~discord.app_commands.CommandNotFound`