"""Decorators that cache results of methods."""
import ast
from functools import wraps
import inspect
from operator import attrgetter
import textwrap
from typing import Any, Callable
import weakref
import numpy as np
# Cached results per instance id, outside the instances so that these can
# still be copied and pickled
_CACHE: dict[int, dict] = {}
[docs]
def cached_until_changed(*names: str,
result: Callable | None = None) -> Callable:
"""
Cache the result of a method until one of the given attributes changes.
The method is called again when the value of any attribute it reads
differs from the value at the last call, also after an array changed in
place. By default the attributes are found in the source of the method:
every chain `self.a.b` that is read gives `a.b`, without a method call
at its end, and a subscript ends the chain. Attributes that are only
assigned are not checked. Copies of an instance have their own cache.
Parameters
----------
*names : str
Attributes to check, dotted names for nested ones. If given, the
source is not searched, and attributes the method reads but names
does not list are not checked.
result : callable, optional
Function of the cached object that gives the returned value, e.g.
the address of a compiled function that must stay alive. By
default the cached object itself is returned.
Returns
-------
callable
Decorator for a method without arguments.
Notes
-----
- Values are compared as numpy arrays, so the attributes must be
numbers, strings, arrays or tuples of them; other objects are an
error. DataArrays backed by dask are loaded on every call.
- The search in the source does not follow other methods or local
names: for `spline = self.spline` only `spline` is checked. A call
`self.method()` or a use of `self` other than `self.a` is an error
unless names are given.
- Copying and comparing the values costs time on every call, e.g.
milliseconds for arrays with 10^5 elements.
- The instances must support weak references.
- Two threads calling for the first time may both compute; the result
of the last one is kept.
"""
def decorator(method: Callable) -> Callable:
getters = [attrgetter(name)
for name in names or _read_attributes(method)]
@wraps(method)
def wrapper(self) -> Any:
values = tuple(_frozen(get(self)) for get in getters)
cache = _instance_cache(self)
cached = cache.get(method.__qualname__)
if cached is None or not _equal(cached[0], values):
cached = (values, method(self))
cache[method.__qualname__] = cached
return cached[1] if result is None else result(cached[1])
return wrapper
return decorator
def _read_attributes(method: Callable) -> tuple[str, ...]:
"""Return the dotted names of the attributes of self a method reads."""
tree = ast.parse(textwrap.dedent(inspect.getsource(method)))
instance = tree.body[0].args.args[0].arg
called: set[int] = {id(node.func) for node in ast.walk(tree)
if isinstance(node, ast.Call)}
bases: set[int] = {id(node.value) for node in ast.walk(tree)
if isinstance(node, ast.Attribute)}
chains: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Name) and node.id == instance:
assert id(node) in bases, (
f"Method {method.__qualname__} uses {instance} other than "
f"in {instance}.a, so the attributes it reads must be given "
"as names.")
if not (isinstance(node, ast.Attribute)
and isinstance(node.ctx, ast.Load)):
continue
parts: list[str] = []
current = node
while isinstance(current, ast.Attribute):
parts.insert(0, current.attr)
current = current.value
if not (isinstance(current, ast.Name) and current.id == instance):
continue
if id(node) in called:
assert len(parts) > 1, (
f"Method {method.__qualname__} calls {instance}.{parts[0]}(), "
"so the attributes it reads must be given as names.")
parts.pop()
chains.add(".".join(parts))
# A chain a.b already covers a
return tuple(sorted(chain for chain in chains
if not any(other.startswith(chain + ".")
for other in chains)))
def _instance_cache(instance: Any) -> dict:
"""Return the cache of an instance, removed when it is deleted."""
key = id(instance)
if key not in _CACHE:
_CACHE[key] = {}
weakref.finalize(instance, _CACHE.pop, key, None)
return _CACHE[key]
def _frozen(value: Any) -> Any:
"""Return a copy of a value, as numpy arrays inside tuples."""
if isinstance(value, (tuple, list)):
return tuple(_frozen(item) for item in value)
frozen = np.array(value, copy=True)
assert frozen.dtype != object, (
f"A value of type {type(value).__name__} cannot be compared, give "
"the numbers inside it as names.")
return frozen
def _equal(old: Any, new: Any) -> bool:
"""Check if two frozen values are equal, NaN equal to NaN."""
if isinstance(old, tuple):
return (isinstance(new, tuple) and len(old) == len(new)
and all(_equal(a, b) for a, b in zip(old, new)))
return (isinstance(new, np.ndarray) and old.shape == new.shape
and old.dtype == new.dtype
and np.array_equal(old, new, equal_nan=old.dtype.kind in "fc"))
