[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):
|
if not isinstance(self.type, BucketType):
|
||||||
raise TypeError('Cooldown type must be a 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()
|
current = time.time()
|
||||||
self._last = current
|
self._last = current
|
||||||
|
|
||||||
|
self._tokens = self.get_tokens(current)
|
||||||
|
|
||||||
# first token used means that we start a new rate limit window
|
# first token used means that we start a new rate limit window
|
||||||
if self._tokens == self.rate:
|
if self._tokens == self.rate:
|
||||||
self._window = current
|
self._window = current
|
||||||
|
|
||||||
# check if our window has passed and we can refresh our tokens
|
# check if we are rate limited
|
||||||
if current > self._window + self.per:
|
|
||||||
self._tokens = self.rate
|
|
||||||
self._window = current
|
|
||||||
|
|
||||||
# check if we're rate limited
|
|
||||||
if self._tokens == 0:
|
if self._tokens == 0:
|
||||||
return self.per - (current - self._window)
|
return self.per - (current - self._window)
|
||||||
|
|
||||||
|
@ -439,13 +439,32 @@ class Command:
|
|||||||
|
|
||||||
if self._buckets.valid:
|
if self._buckets.valid:
|
||||||
bucket = self._buckets.get_bucket(ctx)
|
bucket = self._buckets.get_bucket(ctx)
|
||||||
retry_after = bucket.is_rate_limited()
|
retry_after = bucket.update_rate_limit()
|
||||||
if retry_after:
|
if retry_after:
|
||||||
raise CommandOnCooldown(bucket, retry_after)
|
raise CommandOnCooldown(bucket, retry_after)
|
||||||
|
|
||||||
yield from self._parse_arguments(ctx)
|
yield from self._parse_arguments(ctx)
|
||||||
yield from self.call_before_hooks(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):
|
def reset_cooldown(self, ctx):
|
||||||
"""Resets the cooldown on this command.
|
"""Resets the cooldown on this command.
|
||||||
|
|
||||||
|
Loading…
x
Reference in New Issue
Block a user