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

[commands] Allow passing current to more cooldown mapping methods.

Browse Source 
Also adds a CooldownMapping.update_rate_limit helper function.
This commit is contained in:
Rapptz
2019-04-24 23:23:51 -04:00
parent d80d4145b1
commit 6dcd68b8d7
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -128,20 +128,20 @@ class CooldownMapping:
elif bucket_type is BucketType.category: elif bucket_type is BucketType.category:
return (msg.channel.category or msg.channel).id return (msg.channel.category or msg.channel).id
def _verify_cache_integrity(self): 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
# in a cooldown window. e.g. if we have a command that has a # in a cooldown window. e.g. if we have a command that has a
# cooldown of 60s and it has not been used in 60s then that key should be deleted # cooldown of 60s and it has not been used in 60s then that key should be deleted
current = time.time() current = current or time.time()
dead_keys = [k for k, v in self._cache.items() if current > v._last + v.per] dead_keys = [k for k, v in self._cache.items() if current > v._last + v.per]
for k in dead_keys: for k in dead_keys:
del self._cache[k] del self._cache[k]
def get_bucket(self, message): def get_bucket(self, message, current=None):
if self._cooldown.type is BucketType.default: if self._cooldown.type is BucketType.default:
return self._cooldown return self._cooldown
self._verify_cache_integrity() self._verify_cache_integrity(current)
key = self._bucket_key(message) key = self._bucket_key(message)
if key not in self._cache: if key not in self._cache:
bucket = self._cooldown.copy() bucket = self._cooldown.copy()
@@ -150,3 +150,7 @@ class CooldownMapping:
bucket = self._cache[key] bucket = self._cache[key]
return bucket return bucket
def update_rate_limit(self, message, current=None):
bucket = self.get_bucket(message, current)
return bucket.update_rate_limit(current)

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

@@ -664,8 +664,8 @@ class Command(_BaseCommand):
def _prepare_cooldowns(self, ctx): def _prepare_cooldowns(self, ctx):
if self._buckets.valid: if self._buckets.valid:
bucket = self._buckets.get_bucket(ctx.message)
current = ctx.message.created_at.replace(tzinfo=datetime.timezone.utc).timestamp() current = ctx.message.created_at.replace(tzinfo=datetime.timezone.utc).timestamp()
bucket = self._buckets.get_bucket(ctx.message, current)
retry_after = bucket.update_rate_limit(current) retry_after = bucket.update_rate_limit(current)
if retry_after: if retry_after:
raise CommandOnCooldown(bucket, retry_after) raise CommandOnCooldown(bucket, retry_after)