Add dependable caching

Author: florimondmanca

aiodine/__init__.py

@@ -1,5 +1,7 @@
-from .dependencies import call_resolved, depends
+from .dependencies import CACHE, call_resolved, depends
 
 __version__ = "1.2.8"
 
-__all__ = ["__version__", "depends", "call_resolved"]
+cached = CACHE.cached
+
+__all__ = ["__version__", "depends", "CACHE", "cached", "call_resolved"]

aiodine/dependencies.py

@@ -8,24 +8,14 @@
     asyncnullcontext,
     is_async_context_manager,
 )
+from .models import Dependable, DependableFunc, DependablesCache, T
 
-T = typing.TypeVar("T")
-DependableFunc = typing.Union[
-    typing.Callable[..., typing.Awaitable[T]],
-    typing.Callable[..., typing.AsyncContextManager[T]],
-]
 
+def depends(func: DependableFunc[T], *, cached: bool = None) -> T:
+    return typing.cast(T, Dependable(func, cached=cached))
 
-def depends(func: DependableFunc[T]) -> T:
-    return typing.cast(T, Dependable(func))
 
-
-class Dependable(typing.Generic[T]):
-    def __init__(self, func: DependableFunc[T]):
-        self.func = func
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(func={self.func!r})"
+CACHE = DependablesCache()
 
 
 async def call_resolved(
@@ -46,13 +36,18 @@ async def call_resolved(
     bound = signature.bind_partial(*args, **kwargs)
     bound.apply_defaults()
 
-    for name, value in bound.arguments.items():
-        if isinstance(value, Dependable):
-            bound.arguments[name] = await call_resolved(
-                value.func, __exit_stack__=exit_stack
-            )
+    for name, val in bound.arguments.items():
+        if isinstance(val, Dependable):
+            dependable = val
 
-    raw = func(*bound.args, **bound.kwargs)
+            try:
+                value = CACHE[dependable]
+            except KeyError:
+                value = await call_resolved(dependable.func, __exit_stack__=exit_stack)
+                if CACHE.should_cache(dependable):
+                    CACHE[dependable] = value
+
+            bound.arguments[name] = value
 
     ctx = (
         exit_stack
@@ -61,6 +56,8 @@ async def call_resolved(
     )
 
     async with ctx:
+        raw = func(*bound.args, **bound.kwargs)
+
         if isinstance(raw, types.CoroutineType):
             return await raw
 

aiodine/models.py

@@ -0,0 +1,69 @@
+import typing
+
+T = typing.TypeVar("T")
+DependableFunc = typing.Union[
+    typing.Callable[..., typing.Awaitable[T]],
+    typing.Callable[..., typing.AsyncContextManager[T]],
+]
+F = typing.TypeVar("F", bound=DependableFunc)
+
+
+class Dependable(typing.Generic[T]):
+    __slots__ = ("func", "cached")
+
+    def __init__(self, func: DependableFunc[T], cached: bool = None) -> None:
+        self.func = func
+        self.cached = cached
+
+    def __eq__(self, other: typing.Any) -> bool:
+        if not isinstance(other, Dependable):
+            return False
+        return self.func == other.func and self.cached == other.cached
+
+    def __hash__(self) -> int:
+        return hash((self.func, self.cached))
+
+    def __repr__(self) -> str:
+        attrs = [f"func={self.func!r}"]
+        if self.cached:
+            attrs.append("cached")
+        return f"{self.__class__.__name__}({', '.join(attrs)})"
+
+
+class DependablesCache(typing.Mapping[Dependable, typing.Any]):
+    def __init__(self) -> None:
+        self._last_id = 0
+        self._cached_dependables: typing.Dict[Dependable, typing.Any] = {}
+        self._cachable_funcs: typing.Set[DependableFunc] = set()
+
+    def __getitem__(self, dep: Dependable[T]) -> T:
+        return self._cached_dependables[dep]
+
+    def __setitem__(self, dep: Dependable[T], value: T) -> None:
+        self._cached_dependables[dep] = value
+
+    def __len__(self) -> int:
+        return len(self._cached_dependables)
+
+    def __iter__(self) -> typing.Iterator[Dependable]:
+        return iter(self._cached_dependables)
+
+    def should_cache(self, dep: Dependable[T]) -> bool:
+        if dep.cached:
+            return dep not in self
+        if dep.cached is None:
+            # 'cached=...' was not passed to 'depends()'.
+            # => Should cache if the dependable function was decorated with '@cached'
+            return dep.func in self._cachable_funcs
+        return False
+
+    def cached(self, func: F) -> F:
+        self._cachable_funcs.add(func)
+        return func
+
+    def clear(self) -> None:
+        self._cached_dependables = {}
+        self._cachable_funcs = set()
+
+
+cache = DependablesCache()

tests/models/test_dependable.py

@@ -25,6 +25,17 @@ def __repr__(self) -> str:
         (CowSay(), "Dependable(func=Cow says moo!)"),
     ],
 )
-def test_dependable_repr(func: typing.Callable, output: str) -> None:
-    dependable = aiodine.depends(func)
+@pytest.mark.parametrize("cached", (None, False, True))
+def test_dependable_repr(func: typing.Callable, cached: bool, output: str) -> None:
+    kwargs = {"cached": cached} if cached is not None else {}
+    dependable = aiodine.depends(func, **kwargs)
+    if cached:
+        output = f"{output.rstrip(')')}, cached)"
     assert repr(dependable) == output
+
+
+def test_dependable_equal() -> None:
+    dependable = aiodine.depends(cowsay)
+    assert dependable == dependable
+    assert dependable == aiodine.depends(cowsay)
+    assert dependable != "not a Dependable"

tests/models/test_dependables_cache.py

@@ -0,0 +1,32 @@
+import pytest
+
+import aiodine
+from aiodine.models import DependablesCache
+
+
+@pytest.mark.anyio
+async def test_dependables_cache() -> None:
+    cache = DependablesCache()
+
+    @cache.cached
+    async def moo() -> str:
+        return "moo"
+
+    moo_dependable = aiodine.dependencies.Dependable(moo)
+
+    assert len(cache) == 0
+    assert moo_dependable not in cache
+    with pytest.raises(KeyError):
+        cache[moo_dependable]
+
+    cache[moo_dependable] = await moo()
+
+    assert len(cache) == 1
+    assert moo_dependable in cache
+    assert moo not in cache
+    assert list(cache) == [moo_dependable]
+    assert cache[moo_dependable] == "moo"
+
+    cache.clear()
+    assert len(cache) == 0
+    assert moo_dependable not in cache

tests/test_caching.py

@@ -0,0 +1,55 @@
+import typing
+
+import anyio
+import pytest
+
+import aiodine
+from aiodine.compat import asynccontextmanager
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("cached", (False, True))
+@pytest.mark.parametrize("forced_cached", (None, False, True))
+@pytest.mark.parametrize("use_context_manager", (False, True))
+async def test_cache_function(
+    cached: bool, forced_cached: typing.Optional[bool], use_context_manager: bool
+) -> None:
+    aiodine.CACHE.clear()
+    count = 0
+
+    moo: aiodine.models.DependableFunc[str]
+
+    if use_context_manager:
+
+        @asynccontextmanager
+        async def moo() -> typing.AsyncIterator[str]:
+            nonlocal count
+            count += 1
+            yield "moo"
+
+    else:
+
+        async def moo() -> str:
+            nonlocal count
+            count += 1
+            return "moo"
+
+    if cached:
+        moo = aiodine.cached(moo)
+
+    kwargs = {"cached": forced_cached} if forced_cached is not None else {}
+
+    async def cowsay(moo: str = aiodine.depends(moo, **kwargs)) -> str:
+        return f"Cow says {moo}"
+
+    assert await aiodine.call_resolved(cowsay) == "Cow says moo"
+    assert count == 1
+
+    assert await aiodine.call_resolved(cowsay) == "Cow says moo"
+
+    if (cached or forced_cached) and forced_cached is not False:
+        assert count == 1
+        assert list(aiodine.CACHE) == [aiodine.depends(moo, **kwargs)]
+    else:
+        assert count == 2
+        assert list(aiodine.CACHE) == []