matthew/enhanced-discord.py
1
0
Fork 0
Code Issues 4 Pull Requests 21 Projects Releases Wiki Activity

[commands] allow arbitrary callables in cooldown

Browse Source
This commit is contained in:
Mikey 2021-03-28 03:31:51 -07:00 committed by GitHub
parent baa1ab058b
commit d61486278f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -66,6 +66,9 @@ class BucketType(Enum):
# recieving a DMChannel or GroupChannel which inherit from PrivateChannel and do # 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 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: class Cooldown:
__slots__ = ('rate', 'per', 'type', '_window', '_tokens', '_last') __slots__ = ('rate', 'per', 'type', '_window', '_tokens', '_last')
@ -78,8 +81,8 @@ class Cooldown:
self._tokens = self.rate self._tokens = self.rate
self._last = 0.0 self._last = 0.0
if not isinstance(self.type, BucketType): if not callable(self.type):
raise TypeError('Cooldown type must be a BucketType') raise TypeError('Cooldown type must be a BucketType or callable')
def get_tokens(self, current=None): def get_tokens(self, current=None):
if not current: if not current:
@ -151,7 +154,7 @@ class CooldownMapping:
return cls(Cooldown(rate, per, type)) return cls(Cooldown(rate, per, type))
def _bucket_key(self, msg): def _bucket_key(self, msg):
return self._cooldown.type.get_key(msg) return self._cooldown.type(msg)
def _verify_cache_integrity(self, current=None): def _verify_cache_integrity(self, current=None):
# we want to delete all cache objects that haven't been used # we want to delete all cache objects that haven't been used

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

@ -1959,8 +1959,11 @@ def cooldown(rate, per, type=BucketType.default):
The number of times a command can be used before triggering a cooldown. The number of times a command can be used before triggering a cooldown.
per: :class:`float` per: :class:`float`
The amount of seconds to wait for a cooldown when it's been triggered. The amount of seconds to wait for a cooldown when it's been triggered.
type: :class:`.BucketType` type: Union[:class:`.BucketType`, Callable[[:class:`.Message`], Any]]
The type of cooldown to have. 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): def decorator(func):