diff options
| -rw-r--r-- | discord/ext/commands/cooldowns.py | 9 | ||||
| -rw-r--r-- | discord/ext/commands/core.py | 7 |
2 files changed, 11 insertions, 5 deletions
diff --git a/discord/ext/commands/cooldowns.py b/discord/ext/commands/cooldowns.py index cd7d67ea..54a53396 100644 --- a/discord/ext/commands/cooldowns.py +++ b/discord/ext/commands/cooldowns.py @@ -66,6 +66,9 @@ class BucketType(Enum): # recieving a DMChannel or GroupChannel which inherit from PrivateChannel and do return (msg.channel if isinstance(msg.channel, PrivateChannel) else msg.author.top_role).id + def __call__(self, msg): + return self.get_key(msg) + class Cooldown: __slots__ = ('rate', 'per', 'type', '_window', '_tokens', '_last') @@ -78,8 +81,8 @@ class Cooldown: self._tokens = self.rate self._last = 0.0 - if not isinstance(self.type, BucketType): - raise TypeError('Cooldown type must be a BucketType') + if not callable(self.type): + raise TypeError('Cooldown type must be a BucketType or callable') def get_tokens(self, current=None): if not current: @@ -151,7 +154,7 @@ class CooldownMapping: return cls(Cooldown(rate, per, type)) def _bucket_key(self, msg): - return self._cooldown.type.get_key(msg) + return self._cooldown.type(msg) def _verify_cache_integrity(self, current=None): # we want to delete all cache objects that haven't been used diff --git a/discord/ext/commands/core.py b/discord/ext/commands/core.py index 40c7eaf8..1c22ec0a 100644 --- a/discord/ext/commands/core.py +++ b/discord/ext/commands/core.py @@ -1959,8 +1959,11 @@ def cooldown(rate, per, type=BucketType.default): 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. - type: :class:`.BucketType` - The type of cooldown to have. + type: Union[:class:`.BucketType`, Callable[[:class:`.Message`], Any]] + The type of cooldown to have. If callable, should return a key for the mapping. + + .. versionchanged:: 1.7 + Callables are now supported for custom bucket types. """ def decorator(func): |