mirror
/
Rapptz.discord.py
mirror of https://github.com/Rapptz/discord.py
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

Add application command cooldown decorators 

Fix #7790
pull/7795/head
Rapptz 3 years ago
parent
74b5e0ceb1
commit
9f83eb6032
8 changed files with 312 additions and 129 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 1
      discord/app_commands/__init__.py
  2. 265
      discord/app_commands/checks.py
  3. 23
      discord/app_commands/errors.py
  4. 115
      discord/ext/commands/cooldowns.py
  5. 6
      discord/ext/commands/core.py
  6. 2
      discord/ext/commands/errors.py
  7. 8
      docs/ext/commands/api.rst
  8. 21
      docs/interactions/api.rst

1
discord/app_commands/__init__.py
View File

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

265
discord/app_commands/checks.py
View File

@ -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)

23
discord/app_commands/errors.py
View File

@ -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.

115
discord/ext/commands/cooldowns.py
View File

@ -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,

6
discord/ext/commands/core.py
View File

@ -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`

2
discord/ext/commands/errors.py
View File

@ -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`

8
docs/ext/commands/api.rst
View File

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

21
docs/interactions/api.rst
View File

@ -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`

Write Preview
Loading…
Cancel
Save