Implement ability to specify different cooldown algorithms to use for…

… a command. Closes #255

docs/source/api-reference.rst

@@ -20,6 +20,7 @@ If you think anything is missing, make a merge request to add it, or contact tho
    api_references/context
    api_references/converters
    api_references/cooldowns
+   api_references/cooldown_algorithms
    api_references/decorators
    api_references/errors
    api_references/events

docs/source/changelogs/v2-changelog.rst

@@ -9,6 +9,15 @@ Below are all the changelogs for the stable versions of hikari-lightbulb (versio
 Version 2.3.1
 =============
 
+**Potentially Breaking**
+
+- :obj:`lightbulb.cooldown_algorithms.CooldownStatus` has been moved from the ``buckets`` module to the
+  ``cooldown_algorithms`` module.
+
+- ``commands_run`` attribute has been removed from the :obj:`lightbulb.buckets.Bucket` class.
+
+**Other Changes**
+
 - Add :obj:`lightbulb.commands.base.CommandLike.nsfw` and ``nsfw`` kwarg in :obj:`lightbulb.decorators.command`
   decorator in order to mark a command as only usable in NSFW channels.
 
@@ -22,6 +31,8 @@ Version 2.3.1
 
 - Improve :obj:`lightbulb.converters.special.MessageConverter` to support conversion of `channelid-messageid` format.
 
+- Implement multiple built-in cooldown algorithms which can be specified when adding a cooldown to a command.
+
 Version 2.3.0
 =============
 

docs/source/extension-libs.rst

@@ -174,6 +174,7 @@ Making and running a button group:
         buttons = ExampleButtons(ctx)
         resp = await ctx.respond(f"Tungsten", components = buttons.build())
         await buttons.run(resp)
+
 ----
 
 More coming soon (hopefully).

lightbulb/__init__.py

@@ -20,6 +20,7 @@
     "ApplicationCommand",
     "ApplicationCommandCreationFailed",
     "ApplicationContext",
+    "BangBangCooldownAlgorithm",
     "BaseConverter",
     "BaseHelpCommand",
     "BooleanConverter",
@@ -41,6 +42,7 @@
     "CommandNotFound",
     "Context",
     "ConverterFailure",
+    "CooldownAlgorithm",
     "CooldownManager",
     "CooldownStatus",
     "DefaultHelpCommand",
@@ -100,6 +102,7 @@
     "SlashGroupMixin",
     "SlashSubCommand",
     "SlashSubGroup",
+    "SlidingWindowCooldownAlgorithm",
     "SnowflakeConverter",
     "SubCommandTrait",
     "TextableGuildChannelConverter",
@@ -159,6 +162,7 @@
 from lightbulb import commands
 from lightbulb import context
 from lightbulb import converters
+from lightbulb import cooldown_algorithms
 from lightbulb import cooldowns
 from lightbulb import decorators
 from lightbulb import errors
@@ -172,6 +176,7 @@
 from lightbulb.commands import *
 from lightbulb.context import *
 from lightbulb.converters import *
+from lightbulb.cooldown_algorithms import *
 from lightbulb.cooldowns import *
 from lightbulb.decorators import *
 from lightbulb.errors import *

lightbulb/buckets.py

@@ -17,23 +17,18 @@
 # along with Lightbulb. If not, see <https://www.gnu.org/licenses/>.
 from __future__ import annotations
 
-__all__ = ["Bucket", "UserBucket", "GuildBucket", "GlobalBucket", "ChannelBucket", "CooldownStatus"]
+__all__ = ["Bucket", "UserBucket", "GuildBucket", "GlobalBucket", "ChannelBucket"]
 
 import abc
-import enum
 import time
 import typing as t
 
+from lightbulb import cooldown_algorithms
+
 if t.TYPE_CHECKING:
     from lightbulb.context import base as ctx_base
 
 
-class CooldownStatus(enum.Enum):
-    INACTIVE = enum.auto()
-    ACTIVE = enum.auto()
-    EXPIRED = enum.auto()
-
-
 class Bucket(abc.ABC):
     """
     Base class that represents a bucket that cooldowns or max concurrency limits can
@@ -46,13 +41,19 @@ class Bucket(abc.ABC):
         max_usages (:obj:`int`): Number of command usages before the cooldown is activated.
     """
 
-    __slots__ = ("length", "usages", "commands_run", "activated", "start_time")
+    __slots__ = ("length", "usages", "cooldown_algorithm", "activated", "start_time")
 
-    def __init__(self, length: float, max_usages: int) -> None:
+    def __init__(
+        self,
+        length: float,
+        max_usages: int,
+        cooldown_algorithm: t.Type[
+            cooldown_algorithms.CooldownAlgorithm
+        ] = cooldown_algorithms.BangBangCooldownAlgorithm,
+    ) -> None:
         self.length = length
         self.usages = max_usages
-        self.commands_run: int = 0
-        """Commands run for this bucket since it was created."""
+        self.cooldown_algorithm = cooldown_algorithm()
         self.activated = False
         self.start_time: t.Optional[float] = None
         """The start time of the bucket cooldown. This is relative to :meth:`time.perf_counter`."""
@@ -71,7 +72,9 @@ def extract_hash(cls, context: ctx_base.Context) -> t.Hashable:
         """
         ...
 
-    def acquire(self) -> CooldownStatus:
+    def acquire(
+        self,
+    ) -> t.Union[cooldown_algorithms.CooldownStatus, t.Coroutine[t.Any, t.Any, cooldown_algorithms.CooldownStatus]]:
         """
         Get the current state of the cooldown and add a command usage if the cooldown is not
         currently active or has expired.
@@ -82,15 +85,11 @@ def acquire(self) -> CooldownStatus:
             :obj:`~CooldownStatus`: The status of the cooldown bucket.
         """
         if self.active:
-            return CooldownStatus.ACTIVE
+            return cooldown_algorithms.CooldownStatus.ACTIVE
         elif self.activated and self.expired:
-            return CooldownStatus.EXPIRED
+            return cooldown_algorithms.CooldownStatus.EXPIRED
 
-        self.commands_run += 1
-        if self.commands_run >= self.usages:
-            self.activated = True
-            self.start_time = time.perf_counter()
-        return CooldownStatus.INACTIVE
+        return self.cooldown_algorithm.evaluate(self)
 
     @property
     def active(self) -> bool:
@@ -108,6 +107,10 @@ def expired(self) -> bool:
             return time.perf_counter() >= (self.start_time + self.length)
         return True
 
+    def activate(self) -> None:
+        self.activated = True
+        self.start_time = time.perf_counter()
+
 
 class GlobalBucket(Bucket):
     """

lightbulb/cooldowns.py

@@ -24,6 +24,7 @@
 import typing as t
 
 from lightbulb import buckets
+from lightbulb import cooldown_algorithms
 from lightbulb import errors
 
 if t.TYPE_CHECKING:
@@ -73,19 +74,24 @@ async def add_cooldown(self, context: ctx_base.Context) -> None:
 
         if cooldown_bucket is not None:
             cooldown_status = cooldown_bucket.acquire()
-            if cooldown_status is buckets.CooldownStatus.ACTIVE:
+            if inspect.iscoroutine(cooldown_status):
+                cooldown_status = await cooldown_status
+
+            if cooldown_status is cooldown_algorithms.CooldownStatus.ACTIVE:
                 # Cooldown has been activated
                 assert cooldown_bucket.start_time is not None
                 raise errors.CommandIsOnCooldown(
                     "This command is on cooldown",
                     retry_after=(cooldown_bucket.start_time + cooldown_bucket.length) - time.perf_counter(),
                 )
-            elif cooldown_status is buckets.CooldownStatus.INACTIVE:
+            elif cooldown_status is cooldown_algorithms.CooldownStatus.INACTIVE:
                 # Cooldown has not yet been activated.
                 return
 
         self.cooldowns[cooldown_hash] = bucket
-        self.cooldowns[cooldown_hash].acquire()
+        maybe_coro = self.cooldowns[cooldown_hash].acquire()
+        if inspect.iscoroutine(maybe_coro):
+            await maybe_coro
 
     async def reset_cooldown(self, context: ctx_base.Context) -> None:
         """

lightbulb/decorators.py

@@ -38,6 +38,7 @@
 from lightbulb import buckets
 from lightbulb import checks as checks_
 from lightbulb import commands
+from lightbulb import cooldown_algorithms
 from lightbulb import cooldowns
 
 if t.TYPE_CHECKING:
@@ -266,7 +267,10 @@ def decorate(c_like: commands.base.CommandLike) -> commands.base.CommandLike:
 
 @t.overload
 def add_cooldown(
-    length: float, uses: int, bucket: t.Type[buckets.Bucket]
+    length: float,
+    uses: int,
+    bucket: t.Type[buckets.Bucket],
+    algorithm: t.Type[cooldown_algorithms.CooldownAlgorithm] = cooldown_algorithms.BangBangCooldownAlgorithm,
 ) -> t.Callable[[commands.base.CommandLike], commands.base.CommandLike]:
     ...
 
@@ -283,6 +287,7 @@ def add_cooldown(
     length: t.Optional[float] = None,
     uses: t.Optional[int] = None,
     bucket: t.Optional[t.Type[buckets.Bucket]] = None,
+    algorithm: t.Type[cooldown_algorithms.CooldownAlgorithm] = cooldown_algorithms.BangBangCooldownAlgorithm,
     *,
     callback: t.Optional[
         t.Callable[[context.base.Context], t.Union[buckets.Bucket, t.Coroutine[t.Any, t.Any, buckets.Bucket]]]
@@ -296,6 +301,8 @@ def add_cooldown(
         length (:obj:`float`): The length in seconds of the cooldown timer.
         uses (:obj:`int`): The number of command invocations before the cooldown will be triggered.
         bucket (Type[:obj:`~.buckets.Bucket`]): The bucket to use for cooldowns.
+        algorithm (Type[:obj:`~.cooldown_algorithms.CooldownAlgorithm`]): The cooldown algorithm to use. Defaults
+            to :obj:`~.cooldown_algorithms.BangBangCooldownAlgorithm`.
 
     Keyword Args:
         callback (Callable[[:obj:`~.context.base.Context], Union[:obj:`~.buckets.Bucket`, Coroutine[Any, Any, :obj:`~.buckets.Bucket`]]]): Callable
@@ -307,10 +314,16 @@ def add_cooldown(
     getter: t.Callable[[context.base.Context], t.Union[buckets.Bucket, t.Coroutine[t.Any, t.Any, buckets.Bucket]]]
     if length is not None and uses is not None and bucket is not None:
 
-        def _get_bucket(_: context.base.Context, b: t.Type[buckets.Bucket], l: float, u: int) -> buckets.Bucket:
-            return b(l, u)
+        def _get_bucket(
+            _: context.base.Context,
+            b: t.Type[buckets.Bucket],
+            l: float,
+            u: int,
+            a: t.Type[cooldown_algorithms.CooldownAlgorithm],
+        ) -> buckets.Bucket:
+            return b(l, u, a)
 
-        getter = functools.partial(_get_bucket, b=bucket, l=length, u=uses)
+        getter = functools.partial(_get_bucket, b=bucket, l=length, u=uses, a=algorithm)
     elif callback is not None:
         getter = callback
     else: