manager: cache the output of generator functions
- ID
1e1c897- date
2026-11-07 10:54:12+00:00- author
Alex Chan <alex@alexwlchan.net>- parent
7583cf2- message
manager: cache the output of generator functions Generator functions are materialised into a list, which is saved into storage, then converted back into an generator when retrieved. To avoid unbounded storage, the size of the materialised list is capped at 10,000 items. Larger iterators are returned but not cached.- changed files
2 files, 158 additions, 10 deletions
Changed files
src/funkycache/manager.py (15490 → 17871)
diff --git a/src/funkycache/manager.py b/src/funkycache/manager.py
index e4b0e51..312f6e3 100644
--- a/src/funkycache/manager.py
+++ b/src/funkycache/manager.py
@@ -9,6 +9,7 @@ from collections.abc import Callable, Sequence
from dataclasses import dataclass
from functools import wraps
import inspect
+import itertools
import json
from pathlib import Path
import typing
@@ -183,7 +184,7 @@ class FunkyCache:
"""
def decorator(func: F) -> F:
- data_type = _return_type(func, self._storage)
+ data_type = _data_type(func, self._storage)
dep_paths = _dep_paths(func, files=[Path(p) for p in (files or [])])
dep_values = values or []
@@ -217,18 +218,41 @@ class FunkyCache:
bound.apply_defaults()
key = bound.arguments
+ is_generator = inspect.isgeneratorfunction(func)
+
# We need to look up the CacheStorage in self._known_caches
# every time, in case it's changed since the last call.
storage = self._known_caches[func_name].storage
try:
- return storage.get(key)
+ cached_val = storage.get(key)
+ if is_generator:
+ return cast(R, iter(cached_val))
+ return cached_val
except KeyError:
pass
except Exception as e:
warnings.warn(f"unexpected error looking up key in cache: {e}")
+ # If the function is a generator, capture up to 10,000 items
+ # in a list and store that. If it returns >10,000 values
+ # (for example, if it's infinite), bypass the cache to avoid
+ # storing an unbounded amount of info.
new_value = func(*args, **kwargs)
+ if is_generator:
+ list_value = []
+ for item in itertools.islice(new_value, 10000):
+ list_value.append(item)
+ try:
+ list_value.append(next(new_value))
+ except StopIteration:
+ new_value = list_value
+ else:
+ warnings.warn(
+ "iterator returned more than 10,000 items; "
+ "output will not be cached"
+ )
+ return cast(R, itertools.chain(iter(list_value), new_value))
try:
storage.set(key, new_value)
@@ -238,7 +262,10 @@ class FunkyCache:
warnings.warn(f"unexpected error inserting value into cache: {e}")
pass
- return new_value
+ if is_generator:
+ return cast(R, iter(new_value))
+
+ return cast(R, new_value)
return wrapper
@@ -335,13 +362,15 @@ class FunkyCache:
known_cache.input_sig = new_input_sig
-def _return_type(func: F, storage: StorageType) -> Any:
+def _data_type(func: F, storage: StorageType) -> Any:
"""
- Look up the return annotation of a function.
+ Based on a function's return annotation, look up the type of the
+ data we're going to cache.
"""
sig = inspect.signature(func)
+ ret_annotation = sig.return_annotation
- if sig.return_annotation == inspect.Signature.empty:
+ if ret_annotation == inspect.Signature.empty:
match storage:
case "memory":
warnings.warn(
@@ -356,7 +385,30 @@ def _return_type(func: F, storage: StorageType) -> Any:
case _: # pragma: no cover
raise TypeError(f"unrecognised storage: {storage!r}")
- return sig.return_annotation
+ # If it's a generator function, we'll cache the materialised list
+ # rather than the generator.
+ if inspect.isgeneratorfunction(func):
+ origin = typing.get_origin(ret_annotation)
+ args = typing.get_args(ret_annotation)
+ if origin is not None and len(args) == 1:
+ return list[args[0]] # type: ignore
+ else:
+ match storage:
+ case "memory":
+ warnings.warn(
+ f"{func.__name__} is a generator with an incomplete "
+ "return annotation; "
+ "a complete annotation is required for non-memory storage"
+ )
+ case "sqlite":
+ raise TypeError(
+ f"{func.__name__} is a generator with an incomplete return "
+ "annotation; annotate it to use SQLite storage"
+ )
+ case _: # pragma: no cover
+ raise TypeError(f"unrecognised storage: {storage!r}")
+
+ return ret_annotation
def _canonical_name(func: F) -> str:
tests/test_manager.py (24324 → 27827)
diff --git a/tests/test_manager.py b/tests/test_manager.py
index af49202..85a12b9 100644
--- a/tests/test_manager.py
+++ b/tests/test_manager.py
@@ -7,6 +7,7 @@ from collections.abc import Callable, Iterator, Mapping, Sequence
from collections import Counter
import contextlib
import inspect
+import itertools
from pathlib import Path
import secrets
import subprocess
@@ -78,6 +79,73 @@ class FunkyCacheTestSuite(ABC):
for _ in range(3):
assert initial("apple") == "A"
+ def test_cache_iterator(self, cache: FunkyCache) -> None:
+ """
+ You can cache a function which returns an iterator.
+ """
+ # Count how many times the underlying function was called
+ calls = Counter[str]()
+
+ @cache.register
+ def numbers(max: int) -> Iterator[int]:
+ """
+ Numbers yields all the numbers between 1 and max.
+ """
+ calls[f"max={max}"] += 1
+ for i in itertools.count(start=1):
+ if i > max:
+ break
+ yield i
+ else: # pragma: no cover
+ pass
+
+ # The first call is a cache miss; goes to the real function
+ assert list(numbers(max=3)) == [1, 2, 3]
+ assert calls == {"max=3": 1}
+
+ # The second call is a cache hit; skips the real function
+ assert list(numbers(max=3)) == [1, 2, 3]
+ assert calls == {"max=3": 1}
+
+ # Check that every instance of the iterator is independent
+ n1 = numbers(max=3)
+ n2 = numbers(max=3)
+ assert next(n1) == 1
+ assert next(n1) == 2
+ assert next(n2) == 1
+
+ def test_infinite_iterator_isnt_cached(self, cache: FunkyCache) -> None:
+ """
+ You can register a cache on a function which returns an infinite
+ iterator, but the result won't be cached.
+ """
+ calls = Counter[str]()
+
+ @cache.register
+ def greet_forever(name: str) -> Iterator[str]:
+ calls[name] += 1
+ for i in itertools.count():
+ yield f"Hello {name}! ({i})"
+ else: # pragma: no cover
+ pass
+
+ # The first call is a cache miss
+ with pytest.warns(match="iterator returned more than 10,000 items"):
+ assert next(greet_forever("Alex")) == "Hello Alex! (0)"
+ assert calls == {"Alex": 1}
+
+ # The second call is also a cache miss, because the function was
+ # never cached
+ with pytest.warns(match="iterator returned more than 10,000 items"):
+ assert next(greet_forever("Alex")) == "Hello Alex! (0)"
+ assert calls == {"Alex": 2}
+
+ # The cached items are combined with the rest of the iterator
+ with pytest.warns(match="iterator returned more than 10,000 items"):
+ assert list(itertools.islice(greet_forever("Alex"), 11000)) == [
+ f"Hello Alex! ({i})" for i in range(11000)
+ ]
+
def test_per_function_caches_are_independent(self, cache: FunkyCache) -> None:
"""
Different functions have different caches.
@@ -367,7 +435,8 @@ R = TypeVar("R")
class TestFunkyCache:
"""
Tests for FunkyCache that aren't based on the behaviour of
- a specific storage implementation.
+ a specific storage implementation, but are tricky to test
+ across all implementations.
"""
def test_cache_depends_on_python_file(self, tmp_path: Path) -> None:
@@ -555,7 +624,7 @@ class TestMemoryFunkyCache(FunkyCacheTestSuite):
def test_warns_on_missing_return_signature(self, tmp_path: Path) -> None:
"""
- A function cached with memory storage warns if it doesn't have
+ A function cached in memory warns if it doesn't have
a return signature.
"""
cache = FunkyCache(storage="memory")
@@ -566,6 +635,19 @@ class TestMemoryFunkyCache(FunkyCacheTestSuite):
def double(n: int): # pragma: no cover
return n * 2
+ def test_errors_on_missing_iterator_annotation(self, tmp_path: Path) -> None:
+ """
+ A generator cached in memory warns if a generator doesn't specify
+ the generated type.
+ """
+ cache = FunkyCache(storage="memory")
+
+ with pytest.warns(match="generator with an incomplete return annotation"):
+
+ @cache.register
+ def numbers(n: int) -> Iterator: # pragma: no cover
+ yield from itertools.count(n)
+
class TestSQLiteFunkyCache(FunkyCacheTestSuite):
"""
@@ -675,7 +757,7 @@ class TestSQLiteFunkyCache(FunkyCacheTestSuite):
def test_errors_on_missing_return_signature(self, tmp_path: Path) -> None:
"""
- A function cached with SQLite storage must have a return signature.
+ A function cached in SQLite must have a return signature.
"""
cache = FunkyCache(storage="sqlite", root=tmp_path / ".cache")
@@ -685,6 +767,20 @@ class TestSQLiteFunkyCache(FunkyCacheTestSuite):
def double(n: int): # pragma: no cover
return n * 2
+ def test_errors_on_missing_iterator_annotation(self, tmp_path: Path) -> None:
+ """
+ A generator cached in SQLite must specify the generated type.
+ """
+ cache = FunkyCache(storage="sqlite", root=tmp_path / ".cache")
+
+ with pytest.raises(
+ TypeError, match="generator with an incomplete return annotation"
+ ):
+
+ @cache.register
+ def numbers(n: int) -> Iterator: # pragma: no cover
+ yield from itertools.count(n)
+
@pytest.mark.parametrize("modify_root", path_styles)
@pytest.mark.parametrize("modify_depends_on", path_styles)
@pytest.mark.parametrize(