diff options
Diffstat (limited to 'discord/ext/commands/cooldowns.py')
| -rw-r--r-- | discord/ext/commands/cooldowns.py | 21 |
1 files changed, 14 insertions, 7 deletions
diff --git a/discord/ext/commands/cooldowns.py b/discord/ext/commands/cooldowns.py index 74146284..5849e504 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) |