feat: add the ability to use cooldowns inline via await cooldown.increment()

Author: Skelmis

cooldowns/__init__.py

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

cooldowns/cooldown.py

@@ -245,6 +245,16 @@ def __init__(
         if cooldown_id:
             utils.shared_cooldown_refs[cooldown_id] = self
 
+    async def increment(self, *args, **kwargs) -> "Cooldown":
+        """Inline equivalent to using async with"""
+        if not self._clean_task:
+            self._clean_task = asyncio.create_task(self._keep_buckets_clear())
+
+        last_bucket = await self.get_bucket(*args, **kwargs)
+        bucket: TP = self._get_cooldown_for_bucket(last_bucket)
+        await bucket.increment()
+        return self
+
     async def __aenter__(self) -> "Cooldown":
         if not self._clean_task:
             self._clean_task = asyncio.create_task(self._keep_buckets_clear())
@@ -442,7 +452,9 @@ 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, CallableT]:
+    def bucket(
+        self,
+    ) -> Union[CooldownBucketProtocol, AsyncCooldownBucketProtocol, CallableT]:
         """Returns the underlying bucket to process cooldowns against."""
         return self._bucket
 

cooldowns/cooldown_times_per.py

@@ -46,7 +46,8 @@ def __init__(
     def __repr__(self):
         return f"<CooldownTimesPer(limit={self.limit}, current={self.current}, time_period={self.time_period})>"
 
-    async def __aenter__(self) -> "CooldownTimesPer":
+    async def increment(self) -> "CooldownTimesPer":
+        """Inline equivalent to using async with"""
         if self.current == 0:
             raise CallableOnCooldown(
                 self._cooldown.func, self._cooldown, self.next_reset
@@ -61,6 +62,9 @@ async def __aenter__(self) -> "CooldownTimesPer":
 
         return self
 
+    async def __aenter__(self) -> "CooldownTimesPer":
+        return await self.increment()
+
     async def __aexit__(self, *_) -> None:
         ...
 

cooldowns/protocols/__init__.py

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

cooldowns/protocols/bucket.py

@@ -2,6 +2,7 @@
 
 CallableT = Callable[..., Any] | Coroutine[Any, Any, Any]
 
+
 class CooldownBucketProtocol(Protocol):
     """CooldownBucketProtocol implementation Protocol."""
 

cooldowns/static_times_per.py

@@ -63,7 +63,8 @@ def get_next_reset(self, now: datetime.datetime) -> datetime.datetime:
         possible_options = [self.next_datetime(now, t) for t in self._reset_times]
         return min(possible_options)
 
-    async def __aenter__(self) -> StaticTimesPer:
+    async def increment(self) -> "StaticTimesPer":
+        """Inline equivalent to using async with"""
         if self.current == 0:
             raise CallableOnCooldown(
                 self._cooldown.func, self._cooldown, self.next_reset
@@ -78,3 +79,6 @@ async def __aenter__(self) -> StaticTimesPer:
         self.loop.call_later((reset - now).total_seconds(), self._reset_invoke)
 
         return self
+
+    async def __aenter__(self) -> StaticTimesPer:
+        return await self.increment()

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "function-cooldowns"
-version = "2.1.0"
+version = "2.2.0"
 description = "A simplistic decorator based approach to rate limiting function calls."
 authors = [{name = "Skelmis", email="[email protected]"}]
 license = "UNLICENSE"
@@ -31,6 +31,7 @@ attrs = "^25.1.0"
 pytest = "^8.3.4"
 freezegun = "^1.5.1"
 pytest-asyncio = "^0.25.3"
+black = "^25.1.0"
 
 [build-system]
 requires = ["poetry-core"]

tests/test_cooldown.py

@@ -136,6 +136,7 @@ async def test_func(*args, **kwargs):
 
     await test_func(2)
 
+
 @pytest.mark.asyncio
 async def test_custom_callable_as_bucket():
     def first_arg(*args):
@@ -152,6 +153,7 @@ async def test_func(*args, **kwargs):
 
     await test_func(2)
 
+
 @pytest.mark.asyncio
 async def test_async_custom_callable_as_bucket():
     async def first_arg(*args):
@@ -168,6 +170,7 @@ async def test_func(*args, **kwargs):
 
     await test_func(2)
 
+
 @pytest.mark.asyncio
 async def test_async_bucket_process():
     class CustomBucket(Enum):
@@ -392,3 +395,33 @@ async def test():
     assert _cooldown.get_cooldown_times_per(await _cooldown.get_bucket()) is None
     await test()
     assert _cooldown.get_cooldown_times_per(await _cooldown.get_bucket()) is not None
+
+
+@pytest.mark.asyncio
+async def test_inline_cooldowns():
+    # Can be called once every second
+    # Default bucket is ALL arguments
+    @cooldown(1, 1, bucket=CooldownBucket.all)
+    async def test_func(*args, **kwargs) -> (tuple, dict):
+        return args, kwargs
+
+    _cooldown: Cooldown = getattr(test_func, "_cooldowns")[0]
+    # Call it once, so its on cooldown after this
+    data = await test_func(1, two=2)
+    assert data == ((1,), {"two": 2})
+
+    with pytest.raises(CallableOnCooldown):
+        # Since this uses the same arguments
+        # as the previous call, it comes under
+        # the same bucket, and thus gets rate-limited
+        await _cooldown.increment(1, two=2)
+
+    # Shouldn't error as it comes under the
+    # bucket _HashableArguments(1) rather then
+    # the bucket _HashableArguments(1, two=2)
+    # which are completely different
+    await test_func(1)
+
+    await _cooldown.increment(1, two=3)
+    with pytest.raises(CallableOnCooldown):
+        await _cooldown.increment(1, two=3)