From 9f83eb6032ac2865582076c974f5ff3be1ead962 Mon Sep 17 00:00:00 2001 From: Rapptz Date: Tue, 29 Mar 2022 00:46:42 -0400 Subject: [PATCH] Add application command cooldown decorators Fix #7790 --- discord/app_commands/__init__.py | 1 + discord/app_commands/checks.py | 265 +++++++++++++++++++++++++++++- discord/app_commands/errors.py | 23 +++ discord/ext/commands/cooldowns.py | 115 +------------ discord/ext/commands/core.py | 6 +- discord/ext/commands/errors.py | 2 +- docs/ext/commands/api.rst | 8 - docs/interactions/api.rst | 21 ++- 8 files changed, 312 insertions(+), 129 deletions(-) diff --git a/discord/app_commands/__init__.py b/discord/app_commands/__init__.py index 959a22d81..4a2909862 100644 --- a/discord/app_commands/__init__.py +++ b/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 diff --git a/discord/app_commands/checks.py b/discord/app_commands/checks.py index 762b23c4a..54c88516c 100644 --- a/discord/app_commands/checks.py +++ b/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'' + + 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) diff --git a/discord/app_commands/errors.py b/discord/app_commands/errors.py index 7f14bc3a6..be5c5f601 100644 --- a/discord/app_commands/errors.py +++ b/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. diff --git a/discord/ext/commands/cooldowns.py b/discord/ext/commands/cooldowns.py index 875ef145f..52e125cba 100644 --- a/discord/ext/commands/cooldowns.py +++ b/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'' - - class CooldownMapping: def __init__( self, diff --git a/discord/ext/commands/core.py b/discord/ext/commands/core.py index a2ec67f61..1db6631c0 100644 --- a/discord/ext/commands/core.py +++ b/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 - that cooldown is effectively bypassed. + return a :class:`~discord.app_commands.Cooldown` or ``None``. + 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` diff --git a/discord/ext/commands/errors.py b/discord/ext/commands/errors.py index 2b0567b5b..fc2539658 100644 --- a/discord/ext/commands/errors.py +++ b/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` diff --git a/docs/ext/commands/api.rst b/docs/ext/commands/api.rst index 86671e5b1..ac3357817 100644 --- a/docs/ext/commands/api.rst +++ b/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 -------- diff --git a/docs/interactions/api.rst b/docs/interactions/api.rst index 3513d4c4e..74c577dae 100644 --- a/docs/interactions/api.rst +++ b/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`