feat: allow any callable as a bucket

cooldowns/__init__.py

@@ -51,6 +51,6 @@
     "get_all_cooldowns",
 )
 
-__version__ = "2.0.1"
+__version__ = "2.1.0"
 VersionInfo = namedtuple("VersionInfo", "major minor micro releaselevel serial")
-version_info = VersionInfo(major=2, minor=0, micro=1, releaselevel="final", serial=0)
+version_info = VersionInfo(major=2, minor=1, micro=0, releaselevel="final", serial=0)

cooldowns/cooldown.py

@@ -18,7 +18,7 @@
 )
 from . import CooldownBucket, utils
 from .buckets import _HashableArguments
-from .protocols import CooldownBucketProtocol, AsyncCooldownBucketProtocol
+from .protocols import CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT
 
 logger = getLogger(__name__)
 
@@ -29,7 +29,7 @@
 def cooldown(
     limit: int,
     time_period: Union[float, datetime.timedelta],
-    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol],
+    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT],
     check: Optional[MaybeCoro] = default_check,
     *,
     cooldown_id: Optional[COOLDOWN_ID] = None,
@@ -177,7 +177,7 @@ def __init__(
         limit: int,
         time_period: Union[float, datetime.timedelta],
         bucket: Optional[
-            Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol]
+            Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT]
         ] = None,
         func: Optional[Callable] = None,
         *,
@@ -192,7 +192,7 @@ def __init__(
             period specified by ``time_period``
         time_period: Union[float, datetime.timedelta]
             The time period related to ``limit``. This is seconds.
-        bucket: Optional[Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol]]
+        bucket: Optional[Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT]]
             The :class:`Bucket` implementation to use
             as a bucket to separate cooldown buckets.
 
@@ -230,7 +230,7 @@ def __init__(
 
         self._func: Optional[Callable] = func
         self._bucket: Union[
-            CooldownBucketProtocol, AsyncCooldownBucketProtocol
+            CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT
         ] = bucket
         self.pending_reset: bool = False
         self._raw_last_bucket: dict = {"args": [], "kwargs": {}}
@@ -320,7 +320,12 @@ async def get_bucket(self, *args, **kwargs) -> _HashableArguments:
 
             This can then be used in :meth:`Cooldown.clear` calls.
         """
-        data = await maybe_coro(self._bucket.process, *args, **kwargs)
+        # A cheeky cheat to implement #19
+        # Allows for backwards compat, and makes the assumption if
+        # process doesnt exist then you want to call the bucket itself
+        bucket_method = getattr(self._bucket, "process", self._bucket)
+
+        data = await maybe_coro(bucket_method, *args, **kwargs)
         if self._bucket is CooldownBucket.all:
             return _HashableArguments(*data[0], **data[1])
 
@@ -437,7 +442,7 @@ def __repr__(self) -> str:
         return f"Cooldown(limit={self.limit}, time_period={self.time_period}, func={self._func})"
 
     @property
-    def bucket(self) -> Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol]:
+    def bucket(self) -> Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT]:
         """Returns the underlying bucket to process cooldowns against."""
         return self._bucket
 

cooldowns/protocols/__init__.py

@@ -1,3 +1,3 @@
-from .bucket import CooldownBucketProtocol, AsyncCooldownBucketProtocol
+from .bucket import CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT
 
-__all__ = ("CooldownBucketProtocol", "AsyncCooldownBucketProtocol")
+__all__ = ("CooldownBucketProtocol", "AsyncCooldownBucketProtocol","CallableT")

cooldowns/protocols/bucket.py

@@ -1,5 +1,6 @@
-from typing import Any, Protocol
+from typing import Any, Protocol, Callable, Coroutine
 
+CallableT = Callable[..., Any] | Coroutine[Any, Any, Any]
 
 class CooldownBucketProtocol(Protocol):
     """CooldownBucketProtocol implementation Protocol."""

cooldowns/static_cooldown.py

@@ -7,15 +7,15 @@
 from .buckets import _HashableArguments
 from .cooldown import TP
 from .exceptions import NonExistent
-from .protocols import AsyncCooldownBucketProtocol, CooldownBucketProtocol
+from .protocols import AsyncCooldownBucketProtocol, CooldownBucketProtocol, CallableT
 from .static_times_per import StaticTimesPer
 from .utils import MaybeCoro, default_check, COOLDOWN_ID, maybe_coro
 
 
 def static_cooldown(
     limit: int,
     reset_times: Union[datetime.time, List[datetime.time]],
-    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol],
+    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT],
     check: Optional[MaybeCoro] = default_check,
     *,
     cooldown_id: Optional[COOLDOWN_ID] = None,
@@ -117,7 +117,7 @@ def __init__(
         limit: int,
         reset_times: Union[datetime.time, List[datetime.time]],
         bucket: Optional[
-            Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol]
+            Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT]
         ] = None,
         func: Optional[Callable] = None,
         *,

cooldowns/utils.py

@@ -9,6 +9,7 @@
     NonExistent,
     CooldownAlreadyExists,
 )
+from cooldowns.protocols import CallableT
 
 if TYPE_CHECKING:
     from cooldowns import (
@@ -216,7 +217,7 @@ def get_all_cooldowns(
 def define_shared_cooldown(
     limit: int,
     time_period: float,
-    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol],
+    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT],
     cooldown_id: COOLDOWN_ID,
     *,
     check: Optional[MaybeCoro] = default_check,
@@ -279,7 +280,7 @@ def define_shared_cooldown(
 def define_shared_static_cooldown(
     limit: int,
     reset_times: Union[datetime.time, List[datetime.time]],
-    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol],
+    bucket: Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT],
     cooldown_id: COOLDOWN_ID,
     *,
     check: Optional[MaybeCoro] = default_check,

docs/modules/examples.rst

@@ -177,7 +177,7 @@ Functionally its the same as the previous one.
 Custom buckets
 --------------
 
-All you need is an enum with the ``process`` method.
+All you need is an enum with the ``process`` method or a generic callable.
 
 Heres an example which rate-limits based off of the first argument.
 
@@ -198,6 +198,21 @@ Heres an example which rate-limits based off of the first argument.
     async def test_func(*args, **kwargs):
         .....
 
+You can also do the following:
+
+.. code-block:: python
+    :linenos:
+
+    def first_arg(self, *args, **kwargs):
+        # This bucket is based ONLY off
+        # of the first argument passed
+        return args[0]
+
+    # Then to use
+    @cooldown(1, 1, bucket=first_arg)
+    async def test_func(*args, **kwargs):
+        .....
+
 Example async custom buckets
 ----------------------------
 
@@ -224,6 +239,20 @@ Here is an example which rate-limits based off of the first argument.
     async def test_func(*args, **kwargs):
         .....
 
+You can also do the following:
+
+.. code-block:: python
+    :linenos:
+
+    async def first_arg(self, *args, **kwargs):
+        # This bucket is based ONLY off
+        # of the first argument passed
+        return args[0]
+
+    # Then to use
+    @cooldown(1, 1, bucket=first_arg)
+    async def test_func(*args, **kwargs):
+        .....
 
 Stacking cooldown's
 -------------------

pyproject.toml

@@ -1,19 +1,22 @@
-[tool.poetry]
+[project]
 name = "function-cooldowns"
-version = "2.0.1"
+version = "2.1.0"
 description = "A simplistic decorator based approach to rate limiting function calls."
-authors = ["skelmis <skelmis.craft@gmail.com>"]
+authors = [{name = "Skelmis", email="skelmis.craft@gmail.com"}]
 license = "UNLICENSE"
 readme = "readme.md"
-packages = [{include = "cooldowns"}]
-
-[project]
 classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
 ]
 
+
+[tool.poetry]
+packages = [{include = "cooldowns"}]
+requires-poetry = ">=2.0"
+
+
 [project.urls]
 "Issue tracker" = "https://github.com/Skelmis/Function-Cooldowns/issues"
 Documentation = "https://function-cooldowns.rtfd.io/"

tests/test_cooldown.py

@@ -136,6 +136,37 @@ async def test_func(*args, **kwargs):
 
     await test_func(2)
 
+@pytest.mark.asyncio
+async def test_custom_callable_as_bucket():
+    def first_arg(*args):
+        return args[0]
+
+    @cooldown(1, 1, bucket=first_arg)
+    async def test_func(*args, **kwargs):
+        pass
+
+    await test_func(1, 2, 3)
+
+    with pytest.raises(CallableOnCooldown):
+        await test_func(1)
+
+    await test_func(2)
+
+@pytest.mark.asyncio
+async def test_async_custom_callable_as_bucket():
+    async def first_arg(*args):
+        return args[0]
+
+    @cooldown(1, 1, bucket=first_arg)
+    async def test_func(*args, **kwargs):
+        pass
+
+    await test_func(1, 2, 3)
+
+    with pytest.raises(CallableOnCooldown):
+        await test_func(1)
+
+    await test_func(2)
 
 @pytest.mark.asyncio
 async def test_async_bucket_process():

tests/test_static_cooldown.py

@@ -139,6 +139,29 @@ async def test_func(*args, **kwargs):
     await test_func(2)
 
 
+@pytest.mark.asyncio
+async def test_static_custom_callable_buckets():
+    def first_arg(*args, **kwargs):
+        # This bucket is based ONLY off
+        # of the first argument passed
+        return args[0]
+
+    @static_cooldown(
+        1,
+        datetime.time(hour=3, minute=30, second=1),
+        bucket=first_arg,
+    )
+    async def test_func(*args, **kwargs):
+        pass
+
+    await test_func(1, 2, 3)
+
+    with pytest.raises(CallableOnCooldown):
+        await test_func(1)
+
+    await test_func(2)
+
+
 @pytest.mark.asyncio
 async def test_static_stacked_cooldowns():
     # Can call ONCE time_period second using the same args