Add async lock support and refactor storage classes

Author: zoola969

py_cashier/__init__.py

@@ -1,6 +1,5 @@
 from ._decorators import cache
-from ._key_builders import KeyBuilder
-from ._key_builders._default import DefaultKeyBuilder
+from ._key_builders import DefaultKeyBuilder, KeyBuilder
 from ._serializers import KeySerializer, Md5KeySerializer, ReprKeySerializer, StdHashKeySerializer, StrKeySerializer
 from ._utils import CacheWith
 

py_cashier/_decorators.py

@@ -7,16 +7,16 @@
 from typing_extensions import ParamSpec
 
 from py_cashier._key_builders import DefaultKeyBuilder
-from py_cashier._storages import BaseLock, Result
+from py_cashier._storages import BaseAsyncLock, BaseAsyncStorage, BaseLock, BaseStorage, Result
 from py_cashier.logger import logger
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable
 
     from py_cashier._key_builders import KeyBuilder
-    from py_cashier._storages import BaseStorage
 
 TLock = TypeVar("TLock", bound=BaseLock)
+TAsyncLock = TypeVar("TAsyncLock", bound=BaseAsyncLock)
 P = ParamSpec("P")
 T = TypeVar("T")
 F = TypeVar("F", bound=Callable[..., Any])
@@ -30,24 +30,34 @@ class PStorage(Protocol[T, TLock]):
     def __call__(self) -> BaseStorage[T, TLock]: ...
 
 
+class PAsyncStorage(Protocol[T, TAsyncLock]):
+    def __call__(self) -> BaseAsyncStorage[T, TAsyncLock]: ...
+
+
 def cache(
     *,
-    storage: PStorage[T, TLock],
+    storage: PStorage[T, TLock] | PAsyncStorage[T, TAsyncLock],
     key_builder: PKeyBuilder | None = None,
 ) -> Callable[[F], F]:
     """Cache decorator."""
 
     def _decorator(f: F) -> F:
         k = key_builder() if key_builder is not None else DefaultKeyBuilder(func=f)
+        s = storage()
         if iscoroutinefunction(f):
+            if not isinstance(s, BaseAsyncStorage):
+                msg = "Async function requires an async storage"
+                raise TypeError(msg)
             return cast(
                 "F",
-                _async_wrapper(func=f, storage=storage(), key_builder=k),
+                _async_wrapper(func=f, storage=s, key_builder=k),
             )
-
+        if not isinstance(s, BaseStorage):
+            msg = "Regular function requires a synchronous storage"
+            raise TypeError(msg)
         return cast(
             "F",
-            _wrapper(func=f, storage=storage(), key_builder=k),
+            _wrapper(func=f, storage=s, key_builder=k),
         )
 
     return _decorator
@@ -79,7 +89,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 
 def _async_wrapper(
     func: Callable[P, Awaitable[T]],
-    storage: BaseStorage[T, TLock],
+    storage: BaseAsyncStorage[T, TAsyncLock],
     key_builder: KeyBuilder,
 ) -> Callable[P, Awaitable[T]]:
 

py_cashier/_storages/__init__.py

@@ -1,10 +1,14 @@
-from ._abc import BaseLock, BaseStorage, Result
-from ._ttl_map import SimpleLock, TTLMapStorage
+from ._abc import BaseAsyncLock, BaseAsyncStorage, BaseLock, BaseStorage, Result
+from ._ttl_map import SimpleAsyncLock, SimpleLock, TTLMapAsyncStorage, TTLMapStorage
 
 __all__ = [
+    "BaseAsyncLock",
+    "BaseAsyncStorage",
     "BaseLock",
     "BaseStorage",
     "Result",
+    "SimpleAsyncLock",
     "SimpleLock",
+    "TTLMapAsyncStorage",
     "TTLMapStorage",
 ]

py_cashier/_storages/_abc.py

@@ -30,21 +30,23 @@ def value(self) -> TValue:
 
 class BaseLock(ABC):
     @abstractmethod
-    async def __aenter__(self) -> Self: ...
+    def __enter__(self) -> Self: ...
 
     @abstractmethod
-    async def __aexit__(
+    def __exit__(
         self,
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None: ...
 
+
+class BaseAsyncLock(ABC):
     @abstractmethod
-    def __enter__(self) -> Self: ...
+    async def __aenter__(self) -> Self: ...
 
     @abstractmethod
-    def __exit__(
+    async def __aexit__(
         self,
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
@@ -53,25 +55,32 @@ def __exit__(
 
 
 TLock = TypeVar("TLock", bound=BaseLock)
+TAsyncLock = TypeVar("TAsyncLock", bound=BaseAsyncLock)
 
 
 class BaseStorage(ABC, Generic[TValue, TLock]):
     @abstractmethod
     def lock(self, key: str) -> TLock:
         """Return lock for the key."""
 
-    @abstractmethod
-    async def aget(self, key: str) -> Result[TValue] | None:
-        """Get value by key async."""
-
-    @abstractmethod
-    async def aset(self, key: str, value: TValue) -> None:
-        """Set value by key async."""
-
     @abstractmethod
     def get(self, key: str) -> Result[TValue] | None:
         """Get value by key."""
 
     @abstractmethod
     def set(self, key: str, value: TValue) -> None:
         """Set value by key."""
+
+
+class BaseAsyncStorage(ABC, Generic[TValue, TAsyncLock]):
+    @abstractmethod
+    def lock(self, key: str) -> TAsyncLock:
+        """Return lock for the key."""
+
+    @abstractmethod
+    async def aget(self, key: str) -> Result[TValue] | None:
+        """Get value by key async."""
+
+    @abstractmethod
+    async def aset(self, key: str, value: TValue) -> None:
+        """Set value by key async."""

py_cashier/_storages/_ttl_map.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from asyncio import Condition as AsyncCondition
 from datetime import timedelta
 from functools import partial
 from threading import Condition
@@ -10,7 +11,7 @@
 
 from py_cashier.logger import logger
 
-from ._abc import BaseLock, BaseStorage, Result, TValue
+from ._abc import BaseAsyncLock, BaseAsyncStorage, BaseLock, BaseStorage, Result, TValue
 
 if TYPE_CHECKING:
     from types import TracebackType
@@ -39,6 +40,27 @@ def unregister_lock(self, key: str) -> None:
             self._condition.notify_all()
 
 
+class AsyncLockStorage:
+    def __init__(self) -> None:
+        self._locks: set[str] = set()
+        self._condition = AsyncCondition()
+
+    async def register_lock(self, key: str) -> None:
+        async with self._condition:
+            while key in self._locks:
+                logger.debug("Key '%s' is in use, waiting for release.", key)
+                await self._condition.wait()
+            logger.debug("Registering lock for key '%s'.", key)
+            self._locks.add(key)
+            self._condition.notify_all()
+
+    async def unregister_lock(self, key: str) -> None:
+        async with self._condition:
+            self._locks.discard(key)
+            logger.debug("Unregistering lock for key '%s'.", key)
+            self._condition.notify_all()
+
+
 class SimpleLock(BaseLock):
     def __init__(self, lock_storage: LockStorage, key: str) -> None:
         self._lock_storage = lock_storage
@@ -58,9 +80,15 @@ def __exit__(
     ) -> None:
         self._lock_storage.unregister_lock(self._key)
 
-    # Async context manager methods are useless here, as the lock logic is synchronous.
+
+class SimpleAsyncLock(BaseAsyncLock):
+    def __init__(self, lock_storage: AsyncLockStorage, key: str) -> None:
+        self._lock_storage = lock_storage
+        self._key = key
+
     @override
     async def __aenter__(self) -> Self:
+        await self._lock_storage.register_lock(self._key)
         return self
 
     @override
@@ -70,7 +98,7 @@ async def __aexit__(
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        self._lock_storage.unregister_lock(self._key)
+        await self._lock_storage.unregister_lock(self._key)
 
 
 class TTLMapStorage(BaseStorage[TValue, SimpleLock]):
@@ -101,10 +129,27 @@ def get(self, key: str) -> Result[TValue] | None:
     def set(self, key: str, value: TValue) -> None:
         self._storage[key] = value
 
+
+class TTLMapAsyncStorage(BaseAsyncStorage[TValue, SimpleAsyncLock]):
+    def __init__(
+        self,
+        max_size: int | None = 1024,
+        ttl: timedelta | None = timedelta(minutes=1),
+    ) -> None:
+        self._lock_storage = AsyncLockStorage()
+        self._storage: TTLMap[str, TValue] = TTLMap(max_size=max_size, ttl=ttl)
+
+    @override
+    def lock(self, key: str) -> SimpleAsyncLock:
+        return SimpleAsyncLock(self._lock_storage, key)
+
     @override
     async def aget(self, key: str) -> Result[TValue] | None:
-        return self.get(key)
+        try:
+            return Result(self._storage[key])
+        except KeyError:
+            return None
 
     @override
     async def aset(self, key: str, value: TValue) -> None:
-        return self.set(key, value)
+        self._storage[key] = value

tests/test__decorator.py

@@ -1,13 +1,22 @@
+from __future__ import annotations
+
 import asyncio
 import time
 from collections import defaultdict
 from concurrent.futures.thread import ThreadPoolExecutor
 from datetime import timedelta
+from typing import TYPE_CHECKING, Any
 
 import pytest
 
 from py_cashier import cache
 from py_cashier._storages import TTLMapStorage
+from py_cashier._storages._ttl_map import TTLMapAsyncStorage
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from py_cashier._decorators import PAsyncStorage, PStorage
 
 
 @pytest.mark.parametrize(
@@ -50,7 +59,7 @@ async def test_cache_async(a: int, b: int, expected: int) -> None:
     """Test that the cache decorator works correctly for asynchronous functions."""
     calls: dict[tuple[int, int], int] = defaultdict(int)
 
-    @cache(storage=lambda: TTLMapStorage(max_size=1000, ttl=timedelta(seconds=1)))
+    @cache(storage=lambda: TTLMapAsyncStorage(max_size=1000, ttl=timedelta(seconds=1)))
     async def func(x: int, y: int) -> int:
         nonlocal calls
         calls[(x, y)] += 1
@@ -83,7 +92,7 @@ async def test_cache_dog_piling_async(a: int, b: int, expected: int, tasks_count
     """
     calls = 0
 
-    @cache(storage=lambda: TTLMapStorage(max_size=1000, ttl=timedelta(seconds=1)))
+    @cache(storage=lambda: TTLMapAsyncStorage(max_size=1000, ttl=timedelta(seconds=1)))
     async def func(x: int, y: int) -> int:
         await asyncio.sleep(0.1)
         nonlocal calls
@@ -184,7 +193,7 @@ async def test_cache_failing_func_async(a: int, b: int, tasks_count: int) -> Non
     """Test that failing functions are not cached in asynchronous context."""
     calls = 0
 
-    @cache(storage=lambda: TTLMapStorage(max_size=1000, ttl=timedelta(seconds=1)))
+    @cache(storage=lambda: TTLMapAsyncStorage(max_size=1000, ttl=timedelta(seconds=1)))
     async def func(x: int, y: int) -> int:
         await asyncio.sleep(0.1)
         nonlocal calls
@@ -205,3 +214,24 @@ async def func2(x: int, y: int) -> int:
     assert all(result == 0 for result in results)
     # Each call should execute the function (no caching of errors)
     assert calls == tasks_count
+
+
+async def _af() -> None:
+    pass
+
+
+def _f() -> None:
+    pass
+
+
+@pytest.mark.parametrize(
+    ("func", "storage"),
+    [
+        (_af, lambda: TTLMapStorage(max_size=1000, ttl=timedelta(seconds=1))),
+        (_f, lambda: TTLMapAsyncStorage(max_size=1000, ttl=timedelta(seconds=1))),
+    ],
+)
+def test__decorator__invalid_storage(func: Callable[..., Any], storage: PStorage | PAsyncStorage) -> None:
+    """Test that the decorator raises TypeError for invalid storage."""
+    with pytest.raises(TypeError, match="Regular function requires a synchronous storage"):
+        cache(storage=storage)(func)