Browse Source

[commands] Make CooldownMapping.get_bucket take Message instead.

Requiring a full blown Context might be a bit overkill considering
we only use a single attribute from it.
pull/848/head
Rapptz 8 years ago
parent
commit
1bb7b6ff2d
  1. 7
      discord/ext/commands/cooldowns.py
  2. 6
      discord/ext/commands/core.py

7
discord/ext/commands/cooldowns.py

@ -103,8 +103,7 @@ class CooldownMapping:
def from_cooldown(cls, rate, per, type): def from_cooldown(cls, rate, per, type):
return cls(Cooldown(rate, per, type)) return cls(Cooldown(rate, per, type))
def _bucket_key(self, ctx): def _bucket_key(self, msg):
msg = ctx.message
bucket_type = self._cooldown.type bucket_type = self._cooldown.type
if bucket_type is BucketType.user: if bucket_type is BucketType.user:
return msg.author.id return msg.author.id
@ -122,12 +121,12 @@ class CooldownMapping:
for k in dead_keys: for k in dead_keys:
del self._cache[k] del self._cache[k]
def get_bucket(self, ctx): def get_bucket(self, message):
if self._cooldown.type is BucketType.default: if self._cooldown.type is BucketType.default:
return self._cooldown return self._cooldown
self._verify_cache_integrity() self._verify_cache_integrity()
key = self._bucket_key(ctx) key = self._bucket_key(message)
if key not in self._cache: if key not in self._cache:
bucket = self._cooldown.copy() bucket = self._cooldown.copy()
self._cache[key] = bucket self._cache[key] = bucket

6
discord/ext/commands/core.py

@ -438,7 +438,7 @@ class Command:
yield from self._verify_checks(ctx) yield from self._verify_checks(ctx)
if self._buckets.valid: if self._buckets.valid:
bucket = self._buckets.get_bucket(ctx) bucket = self._buckets.get_bucket(ctx.message)
retry_after = bucket.update_rate_limit() retry_after = bucket.update_rate_limit()
if retry_after: if retry_after:
raise CommandOnCooldown(bucket, retry_after) raise CommandOnCooldown(bucket, retry_after)
@ -462,7 +462,7 @@ class Command:
if not self._buckets.valid: if not self._buckets.valid:
return False return False
bucket = self._buckets.get_bucket(ctx) bucket = self._buckets.get_bucket(ctx.message)
return bucket.get_tokens() == 0 return bucket.get_tokens() == 0
def reset_cooldown(self, ctx): def reset_cooldown(self, ctx):
@ -474,7 +474,7 @@ class Command:
The invocation context to reset the cooldown under. The invocation context to reset the cooldown under.
""" """
if self._buckets.valid: if self._buckets.valid:
bucket = self._buckets.get_bucket(ctx) bucket = self._buckets.get_bucket(ctx.message)
bucket.reset() bucket.reset()
@asyncio.coroutine @asyncio.coroutine

Loading…
Cancel
Save