limits 4.0.1__py3-none-any.whl → 4.1__py3-none-any.whl

limits/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-01-16T11:38:03-0800",
+ "date": "2025-03-07T12:07:00-0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "074be17ab3008f50de700e996d243ba85580b058",
- "version": "4.0.1"
+ "full-revisionid": "aff8ca13e30c9d754690a2931498f023d35dc62f",
+ "version": "4.1"
 }
 '''  # END VERSION_JSON
 

limits/aio/storage/__init__.py

@@ -3,7 +3,7 @@ Implementations of storage backends to be used with
 :class:`limits.aio.strategies.RateLimiter` strategies
 """
 
-from .base import MovingWindowSupport, Storage
+from .base import MovingWindowSupport, SlidingWindowCounterSupport, Storage
 from .etcd import EtcdStorage
 from .memcached import MemcachedStorage
 from .memory import MemoryStorage
@@ -13,6 +13,7 @@ from .redis import RedisClusterStorage, RedisSentinelStorage, RedisStorage
 __all__ = [
     "Storage",
     "MovingWindowSupport",
+    "SlidingWindowCounterSupport",
     "EtcdStorage",
     "MemcachedStorage",
     "MemoryStorage",

limits/aio/storage/base.py

@@ -2,36 +2,35 @@ from __future__ import annotations
 
 import functools
 from abc import ABC, abstractmethod
-from typing import Any, cast
 
 from deprecated.sphinx import versionadded
 
 from limits import errors
 from limits.storage.registry import StorageRegistry
 from limits.typing import (
+    Any,
     Awaitable,
     Callable,
-    List,
     Optional,
     P,
     R,
-    Tuple,
     Type,
     Union,
+    cast,
 )
 from limits.util import LazyDependency
 
 
 def _wrap_errors(
-    storage: Storage,
     fn: Callable[P, Awaitable[R]],
 ) -> Callable[P, Awaitable[R]]:
     @functools.wraps(fn)
     async def inner(*args: P.args, **kwargs: P.kwargs) -> R:  # type: ignore[misc]
+        instance = cast(Storage, args[0])
         try:
             return await fn(*args, **kwargs)
-        except storage.base_exceptions as exc:
-            if storage.wrap_exceptions:
+        except instance.base_exceptions as exc:
+            if instance.wrap_exceptions:
                 raise errors.StorageError(exc) from exc
             raise
 
@@ -44,12 +43,11 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
     Base class to extend when implementing an async storage backend.
     """
 
-    STORAGE_SCHEME: Optional[List[str]]
+    STORAGE_SCHEME: Optional[list[str]]
     """The storage schemes to register against this implementation"""
 
-    def __new__(cls, *args: Any, **kwargs: Any) -> Storage:  # type: ignore[misc]
-        inst = super().__new__(cls)
-
+    def __init_subclass__(cls, **kwargs: Any) -> None:  # type:ignore[explicit-any]
+        super().__init_subclass__(**kwargs)
         for method in {
             "incr",
             "get",
@@ -58,9 +56,8 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
             "reset",
             "clear",
         }:
-            setattr(inst, method, _wrap_errors(inst, getattr(inst, method)))
-
-        return inst
+            setattr(cls, method, _wrap_errors(getattr(cls, method)))
+        super().__init_subclass__(**kwargs)
 
     def __init__(
         self,
@@ -77,7 +74,7 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
 
     @property
     @abstractmethod
-    def base_exceptions(self) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:
+    def base_exceptions(self) -> Union[Type[Exception], tuple[Type[Exception], ...]]:
         raise NotImplementedError
 
     @abstractmethod
@@ -135,24 +132,21 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
 
 class MovingWindowSupport(ABC):
     """
-    Abstract base for storages that intend to support
-    the moving window strategy
+    Abstract base class for async storages that support
+    the :ref:`strategies:moving window` strategy
     """
 
-    def __new__(cls, *args: Any, **kwargs: Any) -> MovingWindowSupport:  # type: ignore[misc]
-        inst = super().__new__(cls)
-
+    def __init_subclass__(cls, **kwargs: Any) -> None:  # type: ignore[explicit-any]
         for method in {
             "acquire_entry",
             "get_moving_window",
         }:
             setattr(
-                inst,
+                cls,
                 method,
-                _wrap_errors(cast(Storage, inst), getattr(inst, method)),
+                _wrap_errors(getattr(cls, method)),
             )
-
-        return inst
+        super().__init_subclass__(**kwargs)
 
     @abstractmethod
     async def acquire_entry(
@@ -169,7 +163,7 @@ class MovingWindowSupport(ABC):
     @abstractmethod
     async def get_moving_window(
         self, key: str, limit: int, expiry: int
-    ) -> Tuple[float, int]:
+    ) -> tuple[float, int]:
         """
         returns the starting point and the number of entries in the moving
         window
@@ -179,3 +173,55 @@ class MovingWindowSupport(ABC):
         :return: (start of window, number of acquired entries)
         """
         raise NotImplementedError
+
+
+class SlidingWindowCounterSupport(ABC):
+    """
+    Abstract base class for async storages that support
+    the :ref:`strategies:sliding window counter` strategy
+    """
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:  # type: ignore[explicit-any]
+        for method in {"acquire_sliding_window_entry", "get_sliding_window"}:
+            setattr(
+                cls,
+                method,
+                _wrap_errors(getattr(cls, method)),
+            )
+        super().__init_subclass__(**kwargs)
+
+    @abstractmethod
+    async def acquire_sliding_window_entry(
+        self,
+        key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        """
+        Acquire an entry if the weighted count of the current and previous
+        windows is less than or equal to the limit
+
+        :param key: rate limit key to acquire an entry in
+        :param limit: amount of entries allowed
+        :param expiry: expiry of the entry
+        :param amount: the number of entries to acquire
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        """
+        Return the previous and current window information.
+
+        :param key: the rate limit key
+        :param expiry: the rate limit expiry, needed to compute the key in some implementations
+        :return: a tuple of (int, float, int, float) with the following information:
+          - previous window counter
+          - previous window TTL
+          - current window counter
+          - current window TTL
+        """
+        raise NotImplementedError

limits/aio/storage/etcd.py

@@ -1,7 +1,7 @@
 import asyncio
 import time
 import urllib.parse
-from typing import TYPE_CHECKING, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from limits.aio.storage.base import Storage
 from limits.errors import ConcurrentUpdateError
@@ -28,6 +28,7 @@ class EtcdStorage(Storage):
         self,
         uri: str,
         max_retries: int = MAX_RETRIES,
+        wrap_exceptions: bool = False,
         **options: str,
     ) -> None:
         """
@@ -35,6 +36,8 @@ class EtcdStorage(Storage):
          ``async+etcd://host:port``,
         :param max_retries: Maximum number of attempts to retry
          in the case of concurrent updates to a rate limit key
+        :param wrap_exceptions: Whether to wrap storage exceptions in
+         :exc:`limits.errors.StorageError` before raising it.
         :param options: all remaining keyword arguments are passed
          directly to the constructor of :class:`aetcd.client.Client`
         :raise ConfigurationError: when :pypi:`aetcd` is not available
@@ -45,11 +48,12 @@ class EtcdStorage(Storage):
             host=parsed.hostname, port=parsed.port, **options
         )
         self.max_retries = max_retries
+        super().__init__(uri, wrap_exceptions=wrap_exceptions)
 
     @property
     def base_exceptions(
         self,
-    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
+    ) -> Union[type[Exception], tuple[type[Exception], ...]]:  # pragma: no cover
         return self.lib.ClientError  # type: ignore[no-any-return]
 
     def prefixed_key(self, key: str) -> bytes:

limits/aio/storage/memcached.py

@@ -1,14 +1,17 @@
 import time
 import urllib.parse
+from collections.abc import Iterable
+from math import ceil, floor
 
 from deprecated.sphinx import versionadded
 
-from limits.aio.storage.base import Storage
-from limits.typing import EmcacheClientP, Optional, Tuple, Type, Union
+from limits.aio.storage.base import SlidingWindowCounterSupport, Storage
+from limits.storage.base import TimestampedSlidingWindow
+from limits.typing import EmcacheClientP, ItemP, Optional, Type, Union
 
 
 @versionadded(version="2.1")
-class MemcachedStorage(Storage):
+class MemcachedStorage(Storage, SlidingWindowCounterSupport, TimestampedSlidingWindow):
     """
     Rate limit storage with memcached as backend.
 
@@ -51,7 +54,7 @@ class MemcachedStorage(Storage):
     @property
     def base_exceptions(
         self,
-    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
+    ) -> Union[Type[Exception], tuple[Type[Exception], ...]]:  # pragma: no cover
         return (
             self.dependency.ClusterNoAvailableNodes,
             self.dependency.CommandError,
@@ -70,19 +73,51 @@ class MemcachedStorage(Storage):
         """
         :param key: the key to get the counter value for
         """
-
         item = await (await self.get_storage()).get(key.encode("utf-8"))
 
         return item and int(item.value) or 0
 
+    async def get_many(self, keys: Iterable[str]) -> dict[bytes, ItemP]:
+        """
+        Return multiple counters at once
+
+        :param keys: the keys to get the counter values for
+        """
+        return await (await self.get_storage()).get_many(
+            [k.encode("utf-8") for k in keys]
+        )
+
     async def clear(self, key: str) -> None:
         """
         :param key: the key to clear rate limits for
         """
         await (await self.get_storage()).delete(key.encode("utf-8"))
 
+    async def decr(self, key: str, amount: int = 1, noreply: bool = False) -> int:
+        """
+        decrements the counter for a given rate limit key
+
+        retursn 0 if the key doesn't exist or if noreply is set to True
+
+        :param key: the key to decrement
+        :param amount: the number to decrement by
+        :param noreply: set to True to ignore the memcached response
+        """
+        storage = await self.get_storage()
+        limit_key = key.encode("utf-8")
+        try:
+            value = await storage.decrement(limit_key, amount, noreply=noreply) or 0
+        except self.dependency.NotFoundCommandError:
+            value = 0
+        return value
+
     async def incr(
-        self, key: str, expiry: int, elastic_expiry: bool = False, amount: int = 1
+        self,
+        key: str,
+        expiry: float,
+        elastic_expiry: bool = False,
+        amount: int = 1,
+        set_expiration_key: bool = True,
     ) -> int:
         """
         increments the counter for a given rate limit key
@@ -92,49 +127,70 @@ class MemcachedStorage(Storage):
         :param elastic_expiry: whether to keep extending the rate limit
          window every hit.
         :param amount: the number to increment by
+        :param set_expiration_key: if set to False, the expiration time won't be stored but the key will still expire
         """
         storage = await self.get_storage()
         limit_key = key.encode("utf-8")
-        expire_key = f"{key}/expires".encode()
-        added = True
+        expire_key = self._expiration_key(key).encode()
+        value = None
         try:
-            await storage.add(limit_key, f"{amount}".encode(), exptime=expiry)
-        except self.dependency.NotStoredStorageCommandError:
-            added = False
-            storage = await self.get_storage()
-
-        if not added:
             value = await storage.increment(limit_key, amount) or amount
-
             if elastic_expiry:
-                await storage.touch(limit_key, exptime=expiry)
-                await storage.set(
-                    expire_key,
-                    str(expiry + time.time()).encode("utf-8"),
-                    exptime=expiry,
-                    noreply=False,
-                )
-
+                await storage.touch(limit_key, exptime=ceil(expiry))
+                if set_expiration_key:
+                    await storage.set(
+                        expire_key,
+                        str(expiry + time.time()).encode("utf-8"),
+                        exptime=ceil(expiry),
+                        noreply=False,
+                    )
+            return value
+        except self.dependency.NotFoundCommandError:
+            # Incrementation failed because the key doesn't exist
+            storage = await self.get_storage()
+            try:
+                await storage.add(limit_key, f"{amount}".encode(), exptime=ceil(expiry))
+                if set_expiration_key:
+                    await storage.set(
+                        expire_key,
+                        str(expiry + time.time()).encode("utf-8"),
+                        exptime=ceil(expiry),
+                        noreply=False,
+                    )
+                value = amount
+            except self.dependency.NotStoredStorageCommandError:
+                # Coult not add the key, probably because a concurrent call has added it
+                storage = await self.get_storage()
+                value = await storage.increment(limit_key, amount) or amount
+                if elastic_expiry:
+                    await storage.touch(limit_key, exptime=ceil(expiry))
+                    if set_expiration_key:
+                        await storage.set(
+                            expire_key,
+                            str(expiry + time.time()).encode("utf-8"),
+                            exptime=ceil(expiry),
+                            noreply=False,
+                        )
             return value
-        else:
-            await storage.set(
-                expire_key,
-                str(expiry + time.time()).encode("utf-8"),
-                exptime=expiry,
-                noreply=False,
-            )
-
-        return amount
 
     async def get_expiry(self, key: str) -> float:
         """
         :param key: the key to get the expiry for
         """
         storage = await self.get_storage()
-        item = await storage.get(f"{key}/expires".encode())
+        item = await storage.get(self._expiration_key(key).encode("utf-8"))
 
         return item and float(item.value) or time.time()
 
+    def _expiration_key(self, key: str) -> str:
+        """
+        Return the expiration key for the given counter key.
+
+        Memcached doesn't natively return the expiration time or TTL for a given key,
+        so we implement the expiration time on a separate key.
+        """
+        return key + "/expires"
+
     async def check(self) -> bool:
         """
         Check if storage is healthy by calling the ``get`` command
@@ -150,3 +206,71 @@ class MemcachedStorage(Storage):
 
     async def reset(self) -> Optional[int]:
         raise NotImplementedError
+
+    async def acquire_sliding_window_entry(
+        self,
+        key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        if amount > limit:
+            return False
+        now = time.time()
+        previous_key, current_key = self.sliding_window_keys(key, expiry, now)
+        (
+            previous_count,
+            previous_ttl,
+            current_count,
+            _,
+        ) = await self._get_sliding_window_info(previous_key, current_key, expiry, now)
+        t0 = time.time()
+        weighted_count = previous_count * previous_ttl / expiry + current_count
+        if floor(weighted_count) + amount > limit:
+            return False
+        else:
+            # Hit, increase the current counter.
+            # If the counter doesn't exist yet, set twice the theorical expiry.
+            # We don't need the expiration key as it is estimated with the timestamps directly.
+            current_count = await self.incr(
+                current_key, 2 * expiry, amount=amount, set_expiration_key=False
+            )
+            t1 = time.time()
+            actualised_previous_ttl = max(0, previous_ttl - (t1 - t0))
+            weighted_count = (
+                previous_count * actualised_previous_ttl / expiry + current_count
+            )
+            if floor(weighted_count) > limit:
+                # Another hit won the race condition: revert the incrementation and refuse this hit
+                # Limitation: during high concurrency at the end of the window,
+                # the counter is shifted and cannot be decremented, so less requests than expected are allowed.
+                await self.decr(current_key, amount, noreply=True)
+                return False
+            return True
+
+    async def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        now = time.time()
+        previous_key, current_key = self.sliding_window_keys(key, expiry, now)
+        return await self._get_sliding_window_info(
+            previous_key, current_key, expiry, now
+        )
+
+    async def _get_sliding_window_info(
+        self, previous_key: str, current_key: str, expiry: int, now: float
+    ) -> tuple[int, float, int, float]:
+        result = await self.get_many([previous_key, current_key])
+
+        raw_previous_count = result.get(previous_key.encode("utf-8"))
+        raw_current_count = result.get(current_key.encode("utf-8"))
+
+        current_count = raw_current_count and int(raw_current_count.value) or 0
+        previous_count = raw_previous_count and int(raw_previous_count.value) or 0
+        if previous_count == 0:
+            previous_ttl = float(0)
+        else:
+            previous_ttl = (1 - (((now - expiry) / expiry) % 1)) * expiry
+        current_ttl = (1 - ((now / expiry) % 1)) * expiry + expiry
+
+        return previous_count, previous_ttl, current_count, current_ttl