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

limits/aio/storage/redis/coredis.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, cast
+
+from limits.aio.storage.redis.bridge import RedisBridge
+from limits.errors import ConfigurationError
+from limits.typing import AsyncCoRedisClient, Callable
+
+if TYPE_CHECKING:
+    import coredis
+
+
+class CoredisBridge(RedisBridge):
+    DEFAULT_CLUSTER_OPTIONS: dict[str, float | str | bool] = {
+        "max_connections": 1000,
+    }
+    "Default options passed to :class:`coredis.RedisCluster`"
+
+    @property
+    def base_exceptions(self) -> type[Exception] | tuple[type[Exception], ...]:
+        return (self.dependency.exceptions.RedisError,)
+
+    def use_sentinel(
+        self,
+        service_name: str | None,
+        use_replicas: bool,
+        sentinel_kwargs: dict[str, str | float | bool] | None,
+        **options: str | float | bool,
+    ) -> None:
+        sentinel_configuration = []
+        connection_options = options.copy()
+
+        sep = self.parsed_uri.netloc.find("@") + 1
+
+        for loc in self.parsed_uri.netloc[sep:].split(","):
+            host, port = loc.split(":")
+            sentinel_configuration.append((host, int(port)))
+        service_name = (
+            self.parsed_uri.path.replace("/", "")
+            if self.parsed_uri.path
+            else service_name
+        )
+
+        if service_name is None:
+            raise ConfigurationError("'service_name' not provided")
+
+        self.sentinel = self.dependency.sentinel.Sentinel(
+            sentinel_configuration,
+            sentinel_kwargs={**self.parsed_auth, **(sentinel_kwargs or {})},
+            **{**self.parsed_auth, **connection_options},
+        )
+        self.storage = self.sentinel.primary_for(service_name)
+        self.storage_replica = self.sentinel.replica_for(service_name)
+        self.connection_getter = lambda readonly: (
+            self.storage_replica if readonly and use_replicas else self.storage
+        )
+
+    def use_basic(self, **options: str | float | bool) -> None:
+        if connection_pool := options.pop("connection_pool", None):
+            self.storage = self.dependency.Redis(
+                connection_pool=connection_pool, **options
+            )
+        else:
+            self.storage = self.dependency.Redis.from_url(self.uri, **options)
+
+        self.connection_getter = lambda _: self.storage
+
+    def use_cluster(self, **options: str | float | bool) -> None:
+        sep = self.parsed_uri.netloc.find("@") + 1
+        cluster_hosts: list[dict[str, int | str]] = []
+        cluster_hosts.extend(
+            {"host": host, "port": int(port)}
+            for loc in self.parsed_uri.netloc[sep:].split(",")
+            if loc
+            for host, port in [loc.split(":")]
+        )
+        self.storage = self.dependency.RedisCluster(
+            startup_nodes=cluster_hosts,
+            **{**self.DEFAULT_CLUSTER_OPTIONS, **self.parsed_auth, **options},
+        )
+        self.connection_getter = lambda _: self.storage
+
+    lua_moving_window: coredis.commands.Script[bytes]
+    lua_acquire_moving_window: coredis.commands.Script[bytes]
+    lua_sliding_window: coredis.commands.Script[bytes]
+    lua_acquire_sliding_window: coredis.commands.Script[bytes]
+    lua_clear_keys: coredis.commands.Script[bytes]
+    lua_incr_expire: coredis.commands.Script[bytes]
+    connection_getter: Callable[[bool], AsyncCoRedisClient]
+
+    def get_connection(self, readonly: bool = False) -> AsyncCoRedisClient:
+        return self.connection_getter(readonly)
+
+    def register_scripts(self) -> None:
+        self.lua_moving_window = self.get_connection().register_script(
+            self.SCRIPT_MOVING_WINDOW
+        )
+        self.lua_acquire_moving_window = self.get_connection().register_script(
+            self.SCRIPT_ACQUIRE_MOVING_WINDOW
+        )
+        self.lua_clear_keys = self.get_connection().register_script(
+            self.SCRIPT_CLEAR_KEYS
+        )
+        self.lua_incr_expire = self.get_connection().register_script(
+            self.SCRIPT_INCR_EXPIRE
+        )
+        self.lua_sliding_window = self.get_connection().register_script(
+            self.SCRIPT_SLIDING_WINDOW
+        )
+        self.lua_acquire_sliding_window = self.get_connection().register_script(
+            self.SCRIPT_ACQUIRE_SLIDING_WINDOW
+        )
+
+    async def incr(
+        self, key: str, expiry: int, elastic_expiry: bool = False, amount: int = 1
+    ) -> int:
+        key = self.prefixed_key(key)
+        value = await self.get_connection().incrby(key, amount)
+        if elastic_expiry or value == amount:
+            await self.get_connection().expire(key, expiry)
+
+        return value
+
+    async def get(self, key: str) -> int:
+        key = self.prefixed_key(key)
+        return int(await self.get_connection(readonly=True).get(key) or 0)
+
+    async def clear(self, key: str) -> None:
+        key = self.prefixed_key(key)
+        await self.get_connection().delete([key])
+
+    async def lua_reset(self) -> int | None:
+        return cast(int, await self.lua_clear_keys.execute([self.prefixed_key("*")]))
+
+    async def get_moving_window(
+        self, key: str, limit: int, expiry: int
+    ) -> tuple[float, int]:
+        key = self.prefixed_key(key)
+        timestamp = time.time()
+        window = await self.lua_moving_window.execute(
+            [key], [timestamp - expiry, limit]
+        )
+        if window:
+            return float(window[0]), window[1]  # type: ignore
+        return timestamp, 0
+
+    async def get_sliding_window(
+        self, previous_key: str, current_key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        previous_key = self.prefixed_key(previous_key)
+        current_key = self.prefixed_key(current_key)
+
+        if window := await self.lua_sliding_window.execute(
+            [previous_key, current_key], [expiry]
+        ):
+            return (
+                int(window[0] or 0),  # type: ignore
+                max(0, float(window[1] or 0)) / 1000,  # type: ignore
+                int(window[2] or 0),  # type: ignore
+                max(0, float(window[3] or 0)) / 1000,  # type: ignore
+            )
+        return 0, 0.0, 0, 0.0
+
+    async def acquire_entry(
+        self, key: str, limit: int, expiry: int, amount: int = 1
+    ) -> bool:
+        key = self.prefixed_key(key)
+        timestamp = time.time()
+        acquired = await self.lua_acquire_moving_window.execute(
+            [key], [timestamp, limit, expiry, amount]
+        )
+
+        return bool(acquired)
+
+    async def acquire_sliding_window_entry(
+        self,
+        previous_key: str,
+        current_key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        previous_key = self.prefixed_key(previous_key)
+        current_key = self.prefixed_key(current_key)
+        acquired = await self.lua_acquire_sliding_window.execute(
+            [previous_key, current_key], [limit, expiry, amount]
+        )
+        return bool(acquired)
+
+    async def get_expiry(self, key: str) -> float:
+        key = self.prefixed_key(key)
+        return max(await self.get_connection().ttl(key), 0) + time.time()
+
+    async def check(self) -> bool:
+        try:
+            await self.get_connection().ping()
+
+            return True
+        except:  # noqa
+            return False
+
+    async def reset(self) -> int | None:
+        prefix = self.prefixed_key("*")
+        keys = await self.storage.keys(prefix)
+        count = 0
+        for key in keys:
+            count += await self.storage.delete([key])
+        return count

limits/aio/storage/redis/redispy.py

@@ -0,0 +1,257 @@
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, cast
+
+from limits.aio.storage.redis.bridge import RedisBridge
+from limits.errors import ConfigurationError
+from limits.typing import AsyncRedisClient, Callable
+
+if TYPE_CHECKING:
+    import redis.commands
+
+
+class RedispyBridge(RedisBridge):
+    DEFAULT_CLUSTER_OPTIONS: dict[str, float | str | bool] = {
+        "max_connections": 1000,
+    }
+    "Default options passed to :class:`redis.asyncio.RedisCluster`"
+
+    @property
+    def base_exceptions(self) -> type[Exception] | tuple[type[Exception], ...]:
+        return (self.dependency.RedisError,)
+
+    def use_sentinel(
+        self,
+        service_name: str | None,
+        use_replicas: bool,
+        sentinel_kwargs: dict[str, str | float | bool] | None,
+        **options: str | float | bool,
+    ) -> None:
+        sentinel_configuration = []
+
+        connection_options = options.copy()
+
+        sep = self.parsed_uri.netloc.find("@") + 1
+
+        for loc in self.parsed_uri.netloc[sep:].split(","):
+            host, port = loc.split(":")
+            sentinel_configuration.append((host, int(port)))
+        service_name = (
+            self.parsed_uri.path.replace("/", "")
+            if self.parsed_uri.path
+            else service_name
+        )
+
+        if service_name is None:
+            raise ConfigurationError("'service_name' not provided")
+
+        self.sentinel = self.dependency.asyncio.Sentinel(
+            sentinel_configuration,
+            sentinel_kwargs={**self.parsed_auth, **(sentinel_kwargs or {})},
+            **{**self.parsed_auth, **connection_options},
+        )
+        self.storage = self.sentinel.master_for(service_name)
+        self.storage_replica = self.sentinel.slave_for(service_name)
+        self.connection_getter = lambda readonly: (
+            self.storage_replica if readonly and use_replicas else self.storage
+        )
+
+    def use_basic(self, **options: str | float | bool) -> None:
+        if connection_pool := options.pop("connection_pool", None):
+            self.storage = self.dependency.asyncio.Redis(
+                connection_pool=connection_pool, **options
+            )
+        else:
+            self.storage = self.dependency.asyncio.Redis.from_url(self.uri, **options)
+
+        self.connection_getter = lambda _: self.storage
+
+    def use_cluster(self, **options: str | float | bool) -> None:
+        sep = self.parsed_uri.netloc.find("@") + 1
+        cluster_hosts = []
+
+        for loc in self.parsed_uri.netloc[sep:].split(","):
+            host, port = loc.split(":")
+            cluster_hosts.append(
+                self.dependency.asyncio.cluster.ClusterNode(host=host, port=int(port))
+            )
+
+        self.storage = self.dependency.asyncio.RedisCluster(
+            startup_nodes=cluster_hosts,
+            **{**self.DEFAULT_CLUSTER_OPTIONS, **self.parsed_auth, **options},
+        )
+        self.connection_getter = lambda _: self.storage
+
+    lua_moving_window: redis.commands.core.Script
+    lua_acquire_moving_window: redis.commands.core.Script
+    lua_sliding_window: redis.commands.core.Script
+    lua_acquire_sliding_window: redis.commands.core.Script
+    lua_clear_keys: redis.commands.core.Script
+    lua_incr_expire: redis.commands.core.Script
+    connection_getter: Callable[[bool], AsyncRedisClient]
+
+    def get_connection(self, readonly: bool = False) -> AsyncRedisClient:
+        return self.connection_getter(readonly)
+
+    def register_scripts(self) -> None:
+        # Redis-py uses a slightly different script registration
+        self.lua_moving_window = self.get_connection().register_script(
+            self.SCRIPT_MOVING_WINDOW
+        )
+        self.lua_acquire_moving_window = self.get_connection().register_script(
+            self.SCRIPT_ACQUIRE_MOVING_WINDOW
+        )
+        self.lua_clear_keys = self.get_connection().register_script(
+            self.SCRIPT_CLEAR_KEYS
+        )
+        self.lua_incr_expire = self.get_connection().register_script(
+            self.SCRIPT_INCR_EXPIRE
+        )
+        self.lua_sliding_window = self.get_connection().register_script(
+            self.SCRIPT_SLIDING_WINDOW
+        )
+        self.lua_acquire_sliding_window = self.get_connection().register_script(
+            self.SCRIPT_ACQUIRE_SLIDING_WINDOW
+        )
+
+    async def incr(
+        self,
+        key: str,
+        expiry: int,
+        elastic_expiry: bool = False,
+        amount: int = 1,
+    ) -> int:
+        """
+        increments the counter for a given rate limit key
+
+
+        :param key: the key to increment
+        :param expiry: amount in seconds for the key to expire in
+        :param amount: the number to increment by
+        """
+        key = self.prefixed_key(key)
+
+        if elastic_expiry:
+            value = await self.get_connection().incrby(key, amount)
+            await self.get_connection().expire(key, expiry)
+            return value
+        else:
+            return cast(int, await self.lua_incr_expire([key], [expiry, amount]))
+
+    async def get(self, key: str) -> int:
+        """
+
+        :param key: the key to get the counter value for
+        """
+
+        key = self.prefixed_key(key)
+        return int(await self.get_connection(readonly=True).get(key) or 0)
+
+    async def clear(self, key: str) -> None:
+        """
+        :param key: the key to clear rate limits for
+
+        """
+        key = self.prefixed_key(key)
+        await self.get_connection().delete(key)
+
+    async def lua_reset(self) -> int | None:
+        return cast(int, await self.lua_clear_keys([self.prefixed_key("*")]))
+
+    async def get_moving_window(
+        self, key: str, limit: int, expiry: int
+    ) -> tuple[float, int]:
+        """
+        returns the starting point and the number of entries in the moving
+        window
+
+        :param key: rate limit key
+        :param expiry: expiry of entry
+        :return: (previous count, previous TTL, current count, current TTL)
+        """
+        key = self.prefixed_key(key)
+        timestamp = time.time()
+        window = await self.lua_moving_window([key], [timestamp - expiry, limit])
+        if window:
+            return float(window[0]), window[1]
+        return timestamp, 0
+
+    async def get_sliding_window(
+        self, previous_key: str, current_key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        if window := await self.lua_sliding_window(
+            [self.prefixed_key(previous_key), self.prefixed_key(current_key)], [expiry]
+        ):
+            return (
+                int(window[0] or 0),
+                max(0, float(window[1] or 0)) / 1000,
+                int(window[2] or 0),
+                max(0, float(window[3] or 0)) / 1000,
+            )
+        return 0, 0.0, 0, 0.0
+
+    async def acquire_entry(
+        self,
+        key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        """
+        :param key: rate limit key to acquire an entry in
+        :param limit: amount of entries allowed
+        :param expiry: expiry of the entry
+
+        """
+        key = self.prefixed_key(key)
+        timestamp = time.time()
+        acquired = await self.lua_acquire_moving_window(
+            [key], [timestamp, limit, expiry, amount]
+        )
+
+        return bool(acquired)
+
+    async def acquire_sliding_window_entry(
+        self,
+        previous_key: str,
+        current_key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        previous_key = self.prefixed_key(previous_key)
+        current_key = self.prefixed_key(current_key)
+        acquired = await self.lua_acquire_sliding_window(
+            [previous_key, current_key], [limit, expiry, amount]
+        )
+        return bool(acquired)
+
+    async def get_expiry(self, key: str) -> float:
+        """
+        :param key: the key to get the expiry for
+        """
+
+        key = self.prefixed_key(key)
+        return max(await self.get_connection().ttl(key), 0) + time.time()
+
+    async def check(self) -> bool:
+        """
+        check if storage is healthy
+        """
+        try:
+            await self.get_connection().ping()
+
+            return True
+        except:  # noqa
+            return False
+
+    async def reset(self) -> int | None:
+        prefix = self.prefixed_key("*")
+        keys = await self.storage.keys(
+            prefix, target_nodes=self.dependency.asyncio.cluster.RedisCluster.ALL_NODES
+        )
+        count = 0
+        for key in keys:
+            count += await self.storage.delete(key)
+        return count

limits/aio/storage/redis/valkey.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from .redispy import RedispyBridge
+
+
+class ValkeyBridge(RedispyBridge):
+    @property
+    def base_exceptions(self) -> type[Exception] | tuple[type[Exception], ...]:
+        return (self.dependency.ValkeyError,)

limits/aio/strategies.py

@@ -2,6 +2,8 @@
 Asynchronous rate limiting strategies
 """
 
+from __future__ import annotations
+
 import time
 from abc import ABC, abstractmethod
 from math import floor, inf
@@ -74,7 +76,7 @@ class MovingWindowRateLimiter(RateLimiter):
         ):
             raise NotImplementedError(
                 "MovingWindowRateLimiting is not implemented for storage "
-                "of type %s" % storage.__class__
+                f"of type {storage.__class__}"
             )
         super().__init__(storage)
 
@@ -200,7 +202,7 @@ class SlidingWindowCounterRateLimiter(RateLimiter):
         ):
             raise NotImplementedError(
                 "SlidingWindowCounterRateLimiting is not implemented for storage "
-                "of type %s" % storage.__class__
+                f"of type {storage.__class__}"
             )
         super().__init__(storage)
 

limits/errors.py

@@ -2,6 +2,8 @@
 errors and exceptions
 """
 
+from __future__ import annotations
+
 
 class ConfigurationError(Exception):
     """

limits/storage/__init__.py

@@ -3,12 +3,14 @@ Implementations of storage backends to be used with
 :class:`limits.strategies.RateLimiter` strategies
 """
 
+from __future__ import annotations
+
 import urllib
 
 import limits  # noqa
 
 from ..errors import ConfigurationError
-from ..typing import Union, cast
+from ..typing import TypeAlias, cast
 from .base import MovingWindowSupport, SlidingWindowCounterSupport, Storage
 from .etcd import EtcdStorage
 from .memcached import MemcachedStorage
@@ -19,11 +21,11 @@ from .redis_cluster import RedisClusterStorage
 from .redis_sentinel import RedisSentinelStorage
 from .registry import SCHEMES
 
-StorageTypes = Union[Storage, "limits.aio.storage.Storage"]
+StorageTypes: TypeAlias = "Storage | limits.aio.storage.Storage"
 
 
 def storage_from_string(
-    storage_string: str, **options: Union[float, str, bool]
+    storage_string: str, **options: float | str | bool
 ) -> StorageTypes:
     """
     Factory function to get an instance of the storage class based
@@ -59,21 +61,22 @@ def storage_from_string(
     scheme = urllib.parse.urlparse(storage_string).scheme
 
     if scheme not in SCHEMES:
-        raise ConfigurationError("unknown storage scheme : %s" % storage_string)
+        raise ConfigurationError(f"unknown storage scheme : {storage_string}")
+
     return cast(StorageTypes, SCHEMES[scheme](storage_string, **options))
 
 
 __all__ = [
-    "storage_from_string",
-    "Storage",
-    "MovingWindowSupport",
-    "SlidingWindowCounterSupport",
     "EtcdStorage",
-    "MongoDBStorageBase",
+    "MemcachedStorage",
     "MemoryStorage",
     "MongoDBStorage",
-    "RedisStorage",
+    "MongoDBStorageBase",
+    "MovingWindowSupport",
     "RedisClusterStorage",
     "RedisSentinelStorage",
-    "MemcachedStorage",
+    "RedisStorage",
+    "SlidingWindowCounterSupport",
+    "Storage",
+    "storage_from_string",
 ]

limits/storage/base.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import functools
-import threading
 from abc import ABC, abstractmethod
 
 from limits import errors
@@ -9,11 +8,8 @@ from limits.storage.registry import StorageRegistry
 from limits.typing import (
     Any,
     Callable,
-    Optional,
     P,
     R,
-    Type,
-    Union,
     cast,
 )
 from limits.util import LazyDependency
@@ -40,7 +36,7 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
     Base class to extend when implementing a storage backend.
     """
 
-    STORAGE_SCHEME: Optional[list[str]]
+    STORAGE_SCHEME: list[str] | None
     """The storage schemes to register against this implementation"""
 
     def __init_subclass__(cls, **kwargs: Any) -> None:  # type: ignore[explicit-any]
@@ -57,22 +53,21 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
 
     def __init__(
         self,
-        uri: Optional[str] = None,
+        uri: str | None = None,
         wrap_exceptions: bool = False,
-        **options: Union[float, str, bool],
+        **options: float | str | bool,
     ):
         """
         :param wrap_exceptions: Whether to wrap storage exceptions in
          :exc:`limits.errors.StorageError` before raising it.
         """
 
-        self.lock = threading.RLock()
         super().__init__()
         self.wrap_exceptions = wrap_exceptions
 
     @property
     @abstractmethod
-    def base_exceptions(self) -> Union[Type[Exception], tuple[Type[Exception], ...]]:
+    def base_exceptions(self) -> type[Exception] | tuple[type[Exception], ...]:
         raise NotImplementedError
 
     @abstractmethod
@@ -112,7 +107,7 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
         raise NotImplementedError
 
     @abstractmethod
-    def reset(self) -> Optional[int]:
+    def reset(self) -> int | None:
         """
         reset storage to clear limits
         """

limits/storage/etcd.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import time
 import urllib.parse
 
 from limits.errors import ConcurrentUpdateError
 from limits.storage.base import Storage
-from limits.typing import TYPE_CHECKING, Optional, Union
+from limits.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     import etcd3
@@ -42,7 +44,7 @@ class EtcdStorage(Storage):
         """
         parsed = urllib.parse.urlparse(uri)
         self.lib = self.dependencies["etcd3"].module
-        self.storage: "etcd3.Etcd3Client" = self.lib.client(
+        self.storage: etcd3.Etcd3Client = self.lib.client(
             parsed.hostname, parsed.port, **options
         )
         self.max_retries = max_retries
@@ -51,7 +53,7 @@ class EtcdStorage(Storage):
     @property
     def base_exceptions(
         self,
-    ) -> Union[type[Exception], tuple[type[Exception], ...]]:  # pragma: no cover
+    ) -> type[Exception] | tuple[type[Exception], ...]:  # pragma: no cover
         return self.lib.Etcd3Exception  # type: ignore[no-any-return]
 
     def prefixed_key(self, key: str) -> bytes:
@@ -127,7 +129,7 @@ class EtcdStorage(Storage):
         except:  # noqa
             return False
 
-    def reset(self) -> Optional[int]:
+    def reset(self) -> int | None:
         return self.storage.delete_prefix(f"{self.PREFIX}/").deleted
 
     def clear(self, key: str) -> None: