diff --git a/discord/ext/commands/cooldowns.py b/discord/ext/commands/cooldowns.py index 741462841..5849e5049 100644 --- a/discord/ext/commands/cooldowns.py +++ b/discord/ext/commands/cooldowns.py @@ -48,20 +48,27 @@ class Cooldown: if not isinstance(self.type, BucketType): raise TypeError('Cooldown type must be a BucketType') - def is_rate_limited(self): + def get_tokens(self, current=None): + if not current: + current = time.time() + + tokens = self._tokens + + if current > self._window + self.per: + tokens = self.rate + return tokens + + def update_rate_limit(self): current = time.time() self._last = current + self._tokens = self.get_tokens(current) + # first token used means that we start a new rate limit window if self._tokens == self.rate: self._window = current - # check if our window has passed and we can refresh our tokens - if current > self._window + self.per: - self._tokens = self.rate - self._window = current - - # check if we're rate limited + # check if we are rate limited if self._tokens == 0: return self.per - (current - self._window) diff --git a/discord/ext/commands/core.py b/discord/ext/commands/core.py index 070e7e3b2..1778ac7f9 100644 --- a/discord/ext/commands/core.py +++ b/discord/ext/commands/core.py @@ -439,13 +439,32 @@ class Command: if self._buckets.valid: bucket = self._buckets.get_bucket(ctx) - retry_after = bucket.is_rate_limited() + retry_after = bucket.update_rate_limit() if retry_after: raise CommandOnCooldown(bucket, retry_after) yield from self._parse_arguments(ctx) yield from self.call_before_hooks(ctx) + def is_on_cooldown(self, ctx): + """Checks whether the command is currently on cooldown. + + Parameters + ----------- + ctx: :class:`.Context.` + The invocation context to use when checking the commands cooldown status. + + Returns + -------- + bool + A boolean indicating if the command is on cooldown. + """ + if not self._buckets.valid: + return False + + bucket = self._buckets.get_bucket(ctx) + return bucket.get_tokens() == 0 + def reset_cooldown(self, ctx): """Resets the cooldown on this command.