[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.
This commit is contained in:
parent
d53cada91e
commit
1bb7b6ff2d
@ -103,8 +103,7 @@ class CooldownMapping:
|
||||
def from_cooldown(cls, rate, per, type):
|
||||
return cls(Cooldown(rate, per, type))
|
||||
|
||||
def _bucket_key(self, ctx):
|
||||
msg = ctx.message
|
||||
def _bucket_key(self, msg):
|
||||
bucket_type = self._cooldown.type
|
||||
if bucket_type is BucketType.user:
|
||||
return msg.author.id
|
||||
@ -122,12 +121,12 @@ class CooldownMapping:
|
||||
for k in dead_keys:
|
||||
del self._cache[k]
|
||||
|
||||
def get_bucket(self, ctx):
|
||||
def get_bucket(self, message):
|
||||
if self._cooldown.type is BucketType.default:
|
||||
return self._cooldown
|
||||
|
||||
self._verify_cache_integrity()
|
||||
key = self._bucket_key(ctx)
|
||||
key = self._bucket_key(message)
|
||||
if key not in self._cache:
|
||||
bucket = self._cooldown.copy()
|
||||
self._cache[key] = bucket
|
||||
|
@ -438,7 +438,7 @@ class Command:
|
||||
yield from self._verify_checks(ctx)
|
||||
|
||||
if self._buckets.valid:
|
||||
bucket = self._buckets.get_bucket(ctx)
|
||||
bucket = self._buckets.get_bucket(ctx.message)
|
||||
retry_after = bucket.update_rate_limit()
|
||||
if retry_after:
|
||||
raise CommandOnCooldown(bucket, retry_after)
|
||||
@ -462,7 +462,7 @@ class Command:
|
||||
if not self._buckets.valid:
|
||||
return False
|
||||
|
||||
bucket = self._buckets.get_bucket(ctx)
|
||||
bucket = self._buckets.get_bucket(ctx.message)
|
||||
return bucket.get_tokens() == 0
|
||||
|
||||
def reset_cooldown(self, ctx):
|
||||
@ -474,7 +474,7 @@ class Command:
|
||||
The invocation context to reset the cooldown under.
|
||||
"""
|
||||
if self._buckets.valid:
|
||||
bucket = self._buckets.get_bucket(ctx)
|
||||
bucket = self._buckets.get_bucket(ctx.message)
|
||||
bucket.reset()
|
||||
|
||||
@asyncio.coroutine
|
||||
|
Loading…
x
Reference in New Issue
Block a user