diff options
| -rw-r--r-- | discord/ext/commands/__init__.py | 1 | ||||
| -rw-r--r-- | discord/ext/commands/cooldowns.py | 121 | ||||
| -rw-r--r-- | discord/ext/commands/core.py | 57 | ||||
| -rw-r--r-- | discord/ext/commands/errors.py | 17 |
4 files changed, 193 insertions, 3 deletions
diff --git a/discord/ext/commands/__init__.py b/discord/ext/commands/__init__.py index ed28b6a0..d3b64a29 100644 --- a/discord/ext/commands/__init__.py +++ b/discord/ext/commands/__init__.py @@ -16,3 +16,4 @@ from .core import * from .errors import * from .formatter import HelpFormatter, Paginator from .converter import * +from .cooldowns import BucketType diff --git a/discord/ext/commands/cooldowns.py b/discord/ext/commands/cooldowns.py new file mode 100644 index 00000000..035ac809 --- /dev/null +++ b/discord/ext/commands/cooldowns.py @@ -0,0 +1,121 @@ +# -*- coding: utf-8 -*- +""" +The MIT License (MIT) + +Copyright (c) 2015-2016 Rapptz + +Permission is hereby granted, free of charge, to any person obtaining a +copy of this software and associated documentation files (the "Software"), +to deal in the Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, sublicense, +and/or sell copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. +""" + +import enum +import time + +__all__ = ['BucketType', 'Cooldown', 'CooldownMapping'] + +class BucketType(enum.Enum): + default = 0 + user = 1 + server = 2 + channel = 3 + +class Cooldown: + __slots__ = ['rate', 'per', 'type', '_window', '_tokens', '_last'] + + def __init__(self, rate, per, type): + self.rate = int(rate) + self.per = float(per) + self.type = type + self._window = 0.0 + self._tokens = self.rate + + if not isinstance(self.type, BucketType): + raise TypeError('Cooldown type must be a BucketType') + + def is_rate_limited(self): + current = time.time() + self._last = 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 + if self._tokens == 0: + return self.per - (current - self._window) + + # we're not so decrement our tokens + self._tokens -= 1 + + # see if we got rate limited due to this token change, and if + # so update the window to point to our current time frame + if self._tokens == 0: + self._window = current + + def copy(self): + return Cooldown(self.rate, self.per, self.type) + + def __repr__(self): + return '<Cooldown rate: {0.rate} per: {0.per} window: {0._window} tokens: {0._tokens}>'.format(self) + +class CooldownMapping: + def __init__(self, original): + self._cache = {} + self._cooldown = original + + @property + def valid(self): + return self._cooldown is not None + + def _bucket_key(self, ctx): + msg = ctx.message + bucket_type = self._cooldown.type + if bucket_type is BucketType.user: + return msg.author.id + elif bucket_type is BucketType.server: + return getattr(msg.server, 'id', msg.author.id) + elif bucket_type is BucketType.channel: + return msg.channel.id + + def _verify_cache_integrity(self): + # we want to delete all cache objects that haven't been used + # in a cooldown window. e.g. if we have a command that has a + # cooldown of 60s and it has not been used in 60s then that key should be deleted + current = time.time() + dead_keys = [k for k, v in self._cache.items() if current > v._last + v.per] + for k in dead_keys: + del self._cache[k] + + def get_bucket(self, ctx): + if self._cooldown.type is BucketType.default: + return self._cooldown + + self._verify_cache_integrity() + key = self._bucket_key(ctx) + if key not in self._cache: + bucket = self._cooldown.copy() + self._cache[key] = bucket + else: + bucket = self._cache[key] + + return bucket diff --git a/discord/ext/commands/core.py b/discord/ext/commands/core.py index fff90ae6..8303181b 100644 --- a/discord/ext/commands/core.py +++ b/discord/ext/commands/core.py @@ -30,12 +30,14 @@ import discord import functools from .errors import * +from .cooldowns import Cooldown, BucketType, CooldownMapping from .view import quoted_word from . import converter as converters __all__ = [ 'Command', 'Group', 'GroupMixin', 'command', 'group', 'has_role', 'has_permissions', 'has_any_role', 'check', - 'bot_has_role', 'bot_has_permissions', 'bot_has_any_role' ] + 'bot_has_role', 'bot_has_permissions', 'bot_has_any_role', + 'cooldown' ] def inject_context(ctx, coro): @functools.wraps(coro) @@ -142,6 +144,7 @@ class Command: self.ignore_extra = kwargs.get('ignore_extra', True) self.instance = None self.parent = None + self._buckets = CooldownMapping(kwargs.get('cooldown')) def dispatch_error(self, error, ctx): try: @@ -328,6 +331,12 @@ class Command: if not self.can_run(ctx): raise CheckFailure('The check functions for command {0.qualified_name} failed.'.format(self)) + if self._buckets.valid: + bucket = self._buckets.get_bucket(ctx) + retry_after = bucket.is_rate_limited() + if retry_after: + raise CommandOnCooldown(bucket, retry_after) + @asyncio.coroutine def invoke(self, ctx): ctx.command = self @@ -637,6 +646,12 @@ def command(name=None, cls=None, **attrs): except AttributeError: checks = [] + try: + cooldown = func.__commands_cooldown__ + del func.__commands_cooldown__ + except AttributeError: + cooldown = None + help_doc = attrs.get('help') if help_doc is not None: help_doc = inspect.cleandoc(help_doc) @@ -647,7 +662,7 @@ def command(name=None, cls=None, **attrs): attrs['help'] = help_doc fname = name or func.__name__ - return cls(name=fname, callback=func, checks=checks, **attrs) + return cls(name=fname, callback=func, checks=checks, cooldown=cooldown, **attrs) return decorator @@ -848,3 +863,41 @@ def bot_has_permissions(**perms): permissions = ch.permissions_for(me) return all(getattr(permissions, perm, None) == value for perm, value in perms.items()) return check(predicate) + +def cooldown(rate, per, type=BucketType.default): + """A decorator that adds a cooldown to a :class:`Command` + or its subclasses. + + A cooldown allows a command to only be used a specific amount + of times in a specific time frame. These cooldowns can be based + either on a per-server, per-channel, per-user, or global basis. + Denoted by the third argument of ``type`` which must be of enum + type ``BucketType`` which could be either: + + - ``BucketType.default`` for a global basis. + - ``BucketType.user`` for a per-user basis. + - ``BucketType.server`` for a per-server basis. + - ``BucketType.channel`` for a per-channel basis. + + If a cooldown is triggered, then :exc:`CommandOnCooldown` is triggered in + :func:`on_command_error` and the local error handler. + + A command can only have a single cooldown. + + Parameters + ------------ + rate: int + The number of times a command can be used before triggering a cooldown. + per: float + The amount of seconds to wait for a cooldown when it's been triggered. + type: ``BucketType`` + The type of cooldown to have. + """ + + def decorator(func): + if isinstance(func, Command): + func.cooldown = Cooldown(rate, per, type) + else: + func.__commands_cooldown__ = Cooldown(rate, per, type) + return func + return decorator diff --git a/discord/ext/commands/errors.py b/discord/ext/commands/errors.py index 1cdb0652..b42b87e9 100644 --- a/discord/ext/commands/errors.py +++ b/discord/ext/commands/errors.py @@ -29,7 +29,7 @@ from discord.errors import DiscordException __all__ = [ 'CommandError', 'MissingRequiredArgument', 'BadArgument', 'NoPrivateMessage', 'CheckFailure', 'CommandNotFound', 'DisabledCommand', 'CommandInvokeError', 'TooManyArguments', - 'UserInputError' ] + 'UserInputError', 'CommandOnCooldown' ] class CommandError(DiscordException): """The base exception type for all command related errors. @@ -110,3 +110,18 @@ class CommandInvokeError(CommandError): self.original = e super().__init__('Command raised an exception: {0.__class__.__name__}: {0}'.format(e)) +class CommandOnCooldown(CommandError): + """Exception raised when the command being invoked is on cooldown. + + Attributes + ----------- + cooldown: Cooldown + A class with attributes ``rate``, ``per``, and ``type`` similar to + the :func:`cooldown` decorator. + retry_after: float + The amount of seconds to wait before you can retry again. + """ + def __init__(self, cooldown, retry_after): + self.cooldown = cooldown + self.retry_after = retry_after + super().__init__('You are on cooldown. Try again in {:.2f}s'.format(retry_after)) |