From 15a2de3789fd0a9c1238dc145044e88a9ce352b5 Mon Sep 17 00:00:00 2001 From: skelmis Date: Sat, 8 Feb 2025 12:07:53 +1300 Subject: [PATCH 1/6] feat: allow straight callables for buckets --- cooldowns/cooldown.py | 7 ++++++- tests/test_cooldown.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/cooldowns/cooldown.py b/cooldowns/cooldown.py index 0ab58fc..cb9b617 100644 --- a/cooldowns/cooldown.py +++ b/cooldowns/cooldown.py @@ -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]) diff --git a/tests/test_cooldown.py b/tests/test_cooldown.py index 3832768..03425b9 100644 --- a/tests/test_cooldown.py +++ b/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(): From e20fb49e3c4dc8ba12c419339d9a8c988e106e8a Mon Sep 17 00:00:00 2001 From: skelmis Date: Sat, 8 Feb 2025 12:09:21 +1300 Subject: [PATCH 2/6] doc: update example docs --- docs/modules/examples.rst | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/docs/modules/examples.rst b/docs/modules/examples.rst index b12d004..a1bf74a 100644 --- a/docs/modules/examples.rst +++ b/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 ------------------- From ef76046de4de971190fcefceadec3d64b2683ea7 Mon Sep 17 00:00:00 2001 From: skelmis Date: Sat, 8 Feb 2025 12:11:57 +1300 Subject: [PATCH 3/6] fix: typings for bucket to be correct --- cooldowns/cooldown.py | 10 +++++----- cooldowns/protocols/__init__.py | 4 ++-- cooldowns/protocols/bucket.py | 3 ++- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/cooldowns/cooldown.py b/cooldowns/cooldown.py index cb9b617..a170c14 100644 --- a/cooldowns/cooldown.py +++ b/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": {}} diff --git a/cooldowns/protocols/__init__.py b/cooldowns/protocols/__init__.py index 7f81027..cae3984 100644 --- a/cooldowns/protocols/__init__.py +++ b/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") diff --git a/cooldowns/protocols/bucket.py b/cooldowns/protocols/bucket.py index 14fea67..096c2ba 100644 --- a/cooldowns/protocols/bucket.py +++ b/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.""" From 8170f0d2628ad4d14335a2dbcefa730846114936 Mon Sep 17 00:00:00 2001 From: skelmis Date: Sat, 8 Feb 2025 12:13:46 +1300 Subject: [PATCH 4/6] chore: bump version --- cooldowns/__init__.py | 4 ++-- pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cooldowns/__init__.py b/cooldowns/__init__.py index cec5a54..2c9f653 100644 --- a/cooldowns/__init__.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index bf1a552..7333579 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "function-cooldowns" -version = "2.0.1" +version = "2.1.0" description = "A simplistic decorator based approach to rate limiting function calls." authors = ["skelmis "] license = "UNLICENSE" From c1e838a31348c2cb53cc184f4ab9ecfeff0f230a Mon Sep 17 00:00:00 2001 From: skelmis Date: Sat, 8 Feb 2025 12:21:51 +1300 Subject: [PATCH 5/6] upgrade to poetry 2.0 syntax --- pyproject.toml | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7333579..f168c06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,19 +1,22 @@ -[tool.poetry] +[project] name = "function-cooldowns" version = "2.1.0" description = "A simplistic decorator based approach to rate limiting function calls." -authors = ["skelmis "] +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/" From f1f3b4cc60370dc757eb6ccaa4184bb06e42eab1 Mon Sep 17 00:00:00 2001 From: skelmis Date: Sat, 8 Feb 2025 12:28:56 +1300 Subject: [PATCH 6/6] feat: add custom callableT support to static cooldowns --- cooldowns/cooldown.py | 2 +- cooldowns/static_cooldown.py | 6 +++--- cooldowns/utils.py | 5 +++-- tests/test_static_cooldown.py | 23 +++++++++++++++++++++++ 4 files changed, 30 insertions(+), 6 deletions(-) diff --git a/cooldowns/cooldown.py b/cooldowns/cooldown.py index a170c14..451baca 100644 --- a/cooldowns/cooldown.py +++ b/cooldowns/cooldown.py @@ -442,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 diff --git a/cooldowns/static_cooldown.py b/cooldowns/static_cooldown.py index f210b8c..1b79a62 100644 --- a/cooldowns/static_cooldown.py +++ b/cooldowns/static_cooldown.py @@ -7,7 +7,7 @@ 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 @@ -15,7 +15,7 @@ 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, *, diff --git a/cooldowns/utils.py b/cooldowns/utils.py index 0446206..bbb1744 100644 --- a/cooldowns/utils.py +++ b/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, diff --git a/tests/test_static_cooldown.py b/tests/test_static_cooldown.py index 84fd971..0590768 100644 --- a/tests/test_static_cooldown.py +++ b/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