[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.
This commit is contained in:
parent
b22f7c76c5
commit
bae6f80327
@ -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)
|
||||
|
||||
|
@ -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…
x
Reference in New Issue
Block a user