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
2 changed files with 11 additions and 5 deletions
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
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