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

Also adds a CooldownMapping.update_rate_limit helper function.

discord/ext/commands/cooldowns.py

@@ -128,20 +128,20 @@ class CooldownMapping:
         elif bucket_type is BucketType.category:
             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
         # 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
-        current = time.time()
+        current = current or time.time()
         dead_keys = [k for k, v in self._cache.items() if current > v._last + v.per]
         for k in dead_keys:
             del self._cache[k]
 
-    def get_bucket(self, message):
+    def get_bucket(self, message, current=None):
         if self._cooldown.type is BucketType.default:
             return self._cooldown
 
-        self._verify_cache_integrity()
+        self._verify_cache_integrity(current)
         key = self._bucket_key(message)
         if key not in self._cache:
             bucket = self._cooldown.copy()
@@ -150,3 +150,7 @@ class CooldownMapping:
             bucket = self._cache[key]
 
         return bucket
+
+    def update_rate_limit(self, message, current=None):
+        bucket = self.get_bucket(message, current)
+        return bucket.update_rate_limit(current)

discord/ext/commands/core.py

@@ -664,8 +664,8 @@ class Command(_BaseCommand):
 
     def _prepare_cooldowns(self, ctx):
         if self._buckets.valid:
-            bucket = self._buckets.get_bucket(ctx.message)
             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)
             if retry_after:
                 raise CommandOnCooldown(bucket, retry_after)