Browse Source

[commands] Fix errors when DynamicCooldown returns None

pull/8048/head
z03h 3 years ago
committed by GitHub
parent
commit
2be65e5874
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
  1. 6
      discord/ext/commands/cooldowns.py
  2. 7
      discord/ext/commands/core.py

6
discord/ext/commands/cooldowns.py

@ -125,9 +125,9 @@ class CooldownMapping:
def create_bucket(self, message: Message) -> Cooldown: def create_bucket(self, message: Message) -> Cooldown:
return self._cooldown.copy() # type: ignore return self._cooldown.copy() # type: ignore
def get_bucket(self, message: Message, current: Optional[float] = None) -> Cooldown: def get_bucket(self, message: Message, current: Optional[float] = None) -> Optional[Cooldown]:
if self._type is BucketType.default: if self._type is BucketType.default:
return self._cooldown # type: ignore return self._cooldown
self._verify_cache_integrity(current) self._verify_cache_integrity(current)
key = self._bucket_key(message) key = self._bucket_key(message)
@ -142,6 +142,8 @@ class CooldownMapping:
def update_rate_limit(self, message: Message, current: Optional[float] = None, tokens: int = 1) -> Optional[float]: def update_rate_limit(self, message: Message, current: Optional[float] = None, tokens: int = 1) -> Optional[float]:
bucket = self.get_bucket(message, current) bucket = self.get_bucket(message, current)
if bucket is None:
return None
return bucket.update_rate_limit(current, tokens=tokens) return bucket.update_rate_limit(current, tokens=tokens)

7
discord/ext/commands/core.py

@ -930,6 +930,8 @@ class Command(_BaseCommand, Generic[CogT, P, T]):
return False return False
bucket = self._buckets.get_bucket(ctx.message) bucket = self._buckets.get_bucket(ctx.message)
if bucket is None:
return False
dt = ctx.message.edited_at or ctx.message.created_at dt = ctx.message.edited_at or ctx.message.created_at
current = dt.replace(tzinfo=datetime.timezone.utc).timestamp() current = dt.replace(tzinfo=datetime.timezone.utc).timestamp()
return bucket.get_tokens(current) == 0 return bucket.get_tokens(current) == 0
@ -948,7 +950,8 @@ class Command(_BaseCommand, Generic[CogT, P, T]):
""" """
if self._buckets.valid: if self._buckets.valid:
bucket = self._buckets.get_bucket(ctx.message) bucket = self._buckets.get_bucket(ctx.message)
bucket.reset() if bucket is not None:
bucket.reset()
def get_cooldown_retry_after(self, ctx: Context[BotT], /) -> float: def get_cooldown_retry_after(self, ctx: Context[BotT], /) -> float:
"""Retrieves the amount of seconds before this command can be tried again. """Retrieves the amount of seconds before this command can be tried again.
@ -972,6 +975,8 @@ class Command(_BaseCommand, Generic[CogT, P, T]):
""" """
if self._buckets.valid: if self._buckets.valid:
bucket = self._buckets.get_bucket(ctx.message) bucket = self._buckets.get_bucket(ctx.message)
if bucket is None:
return 0.0
dt = ctx.message.edited_at or ctx.message.created_at dt = ctx.message.edited_at or ctx.message.created_at
current = dt.replace(tzinfo=datetime.timezone.utc).timestamp() current = dt.replace(tzinfo=datetime.timezone.utc).timestamp()
return bucket.get_retry_after(current) return bucket.get_retry_after(current)

Loading…
Cancel
Save