Browse Source

[commands] Split Cooldown state processing to two different functions.

This allows us to check if we are rate limited without
creating a new cool-down window for the command.
pull/812/merge
MysterialPy 8 years ago
committed by Rapptz
parent
commit
bae6f80327
  1. 21
      discord/ext/commands/cooldowns.py
  2. 21
      discord/ext/commands/core.py

21
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)

21
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.

Loading…
Cancel
Save