mcrotk-cache-utils 0.1.0__py3-none-any.whl

mcrotk_cache_utils/decorators/method_caching.py

@@ -0,0 +1,204 @@
+"""Caching decorators for methods, with a separate cache per instance."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Callable, Iterable, Mapping, MutableMapping
+from typing import TYPE_CHECKING, Any, TypeVar, overload
+
+import cachetools
+
+from ..copying.copy_handling import (
+    CACHE_ATTR,
+    ON_COPY,
+    ON_COPY_VALUES,
+    InstallsCopyMethods,
+    OnCopy,
+    check_copy_handling,
+)
+from ..keys.cache_keys import ArgRef, ArgSpec, Passthrough
+from ..keys.passthrough import make_view
+from .support.decorator_options import MethodOptions
+from .support.placement_checks import MethodInClassBody, add_check, reject_static_or_class_method
+
+
+class _DescriptorChecks(MethodInClassBody, InstallsCopyMethods):
+    """What we add to cachetools' method descriptor: the placement error and copy handling."""
+
+
+if TYPE_CHECKING:
+    from typing_extensions import Unpack
+
+F = TypeVar("F", bound=Callable[..., Any])  # the decorated function: its type is kept as it is
+
+
+@overload
+def cached_method(func: F, /, **options: Unpack[MethodOptions]) -> F: ...
+@overload
+def cached_method(func: None = None, /, **options: Unpack[MethodOptions]) -> Callable[[F], F]: ...
+def cached_method(
+    func: Callable[..., Any] | None = None,
+    /,
+    *,
+    cache_factory: Callable[[], MutableMapping] = dict,
+    omit: ArgRef | Iterable[ArgRef] = (),
+    passthrough: Mapping[ArgRef, int | Passthrough] | None = None,
+    key_converters: Mapping[ArgRef, Callable[[Any], Any]] | None = None,
+    lock: Callable[[Any], Any] | None = None,
+    condition: Callable[[Any], Any] | None = None,
+    info: bool = False,
+    on_copy: OnCopy | None = None,
+) -> Any:
+    """Cache a method's results, with a separate cache for each instance.
+
+    Use as ``@cached_method`` or ``@cached_method(...)``. Built on
+    :func:`cachetools.cachedmethod`, which does the calling, binding, locking and statistics.
+    For ordinary methods only: for functions, static methods and class methods use
+    :func:`cached_function` (using this decorator on them is an error).
+
+    **The cache.** Created on an instance's first call and stored in its ``__dict__`` as
+    ``_<Class>_<method>_cache`` (``<Class>`` is the class the method is defined in, so a
+    decorated override and the parent method it calls with ``super()`` don't share a cache).
+    It is not a dataclass field, so it doesn't affect ``==``, ``repr`` or frozen dataclasses,
+    and it disappears with the instance. The class needs a ``__dict__`` (no ``__slots__``-only
+    classes).
+
+    **The key.** Built from the call's arguments bound to the signature with defaults filled
+    in, so ``f(1)``, ``f(x=1)`` and ``f()`` with default ``x=1`` are the same call. ``self`` is
+    always left out. Instance attributes are not part of the key: if you change an attribute
+    that affects results, call :func:`clear_caches`.
+
+    **Types.** Type checkers see the method's own type, so calls are checked against its
+    signature and overriding works as usual (cached or not, in either class). The cache
+    attributes cachetools adds (``obj.method.cache_info()``, ...) work, but aren't part of that
+    type: :func:`cache_controls` gives typed access to them.
+
+    **Unhashable arguments** (also among ``*args`` and ``**kwargs``) are converted:
+
+    - numpy arrays (``ndarray`` and ``memmap``) by their contents: shape, dtype and a hash of
+      the data, computed on every call. So equal arrays share a result, and an array changed in
+      place gets a new key. Object arrays and other array types (masked arrays, ...) need
+      ``key_converters``.
+    - lists, dicts and sets, and tuples and frozensets containing them, by
+      :func:`container_key` with ``type_tags=False``: lists become tuples, sets frozensets and
+      dicts frozensets of their ``(key, value)`` pairs, recursively. So ``[1, 2]`` and
+      ``(1, 2)`` share a result. If the result depends on which was passed (e.g. it has the
+      same type), keep the types with ``key_converters={'x': container_key}``.
+
+    Anything else unhashable (including pandas objects) needs ``key_converters`` or ``omit``.
+
+    Args:
+        cache_factory: Makes each instance's cache. Default ``dict`` (unbounded). Any
+            cachetools cache works, e.g. ``lambda: cachetools.LRUCache(maxsize=128)``.
+        omit: Arguments to leave out of the key, by name or by position in the signature
+            (positions count ``self``: in ``def f(self, a, b)``, ``a`` is 1; negative positions
+            count from the end). Only omit arguments that don't affect the result: calls that
+            differ only in an omitted argument get the same cached result.
+        passthrough: ``{argument: output_position}`` or ``{argument: (output_position, transform)}``.
+            For methods that take an argument and return it, possibly updated, as part of a
+            tuple (e.g. a random number generator). The argument is left out of the key. On a
+            miss the method's own outputs are returned. On a hit the stored result is returned
+            with the caller's current argument (after ``transform``, if given) at
+            ``output_position``. The method must return a tuple; NamedTuples are rebuilt as the
+            same type.
+        key_converters: ``{argument: function}``: the function's result is used in the key in
+            place of the argument's value, e.g. to make an unhashable argument hashable, or to
+            key an argument by part of it (``{'model': lambda m: m.name}``). It replaces the
+            automatic handling of arrays and containers.
+            Ready-made converters: :func:`container_key` (lists, dicts, sets, nested),
+            :func:`array_key`, and for pandas objects (only when nothing simpler will do)
+            :func:`dataframe_key`, :func:`series_key` and :func:`index_key`.
+        lock, condition, info: Passed to :func:`cachetools.cachedmethod` (``lock`` and
+            ``condition`` are functions of ``self`` returning the instance's lock/condition).
+            With ``info=True``, ``cache_controls(obj.method).cache_info()`` gives hit and
+            miss counts.
+        on_copy: What a copy of the instance (``copy.copy`` or ``copy.deepcopy``) gets in
+            this method's cache: ``'empty'`` (the default: never stale, even if the copy is
+            then changed), ``'copy'`` (the cached results: shared with the original by
+            ``copy.copy``, deep-copied by ``deepcopy``; stale if the copy is changed, until
+            :func:`clear_caches`), or ``'manual'`` (the class's own ``__copy__`` handles it, e.g.
+            with :func:`copy_instance`). For a different choice for one copy, use
+            :func:`copy_instance`.
+
+    **Copying.** The decorator gives the class a ``__copy__`` and ``__deepcopy__`` that do
+    Python's usual copy but apply ``on_copy`` to each cached method's cache, and make the
+    copy's cached methods run on the copy. Without them, ``copy.copy`` would give the copy the
+    original's caches, and its cached methods would run on the original (a cachetools
+    behavior). If the class (or a base) defines its own ``__copy__`` or ``__deepcopy__``,
+    that one is kept: setting ``on_copy`` to ``'empty'`` or ``'copy'`` is then an error, and
+    the default gives a :class:`CacheCopyWarning` (once per class) unless the methods use
+    ``on_copy='manual'``. ``'manual'`` requires a ``__copy__`` of the class's own. Copies
+    start with zeroed ``info`` counters. Pickling is unaffected: it carries the caches along.
+    ``dataclasses.replace`` and ``copy.replace`` make new instances, with empty caches.
+
+    Errors in the options (unknown arguments, conflicts) are raised when the decorator is
+    applied. An unhashable key raises ``TypeError`` naming the argument.
+
+    Caveats:
+        - Results are returned as stored, not copied: changing a returned object in place
+          changes the cached result.
+        - Pickling an instance (e.g. sending it to worker processes) carries its caches along.
+        - Copying a cache for ``on_copy='copy'`` isn't thread-safe: don't copy an instance
+          while another thread is calling its cached methods.
+        - With ``passthrough``, cached entries record the type of the returned tuple. If that's
+          a NamedTuple class defined inside a function, an instance with cached results can't
+          be pickled (the same limitation as pickling any object referring to a local class).
+    """
+
+    if on_copy is not None and on_copy not in ON_COPY_VALUES:
+        raise ValueError(f"cached_method: on_copy must be one of {ON_COPY_VALUES}, got {on_copy!r}")
+
+    def decorate(func: Callable[..., Any]) -> Any:
+        """Wrap ``func`` with cachetools' method descriptor, with its key spec and per-instance caches."""
+        reject_static_or_class_method(
+            func,
+            "cached_method",
+            "For a static or class method, use @cached_function under @staticmethod/@classmethod.",
+        )
+        spec = ArgSpec(func, has_self=True, omit=omit, passthrough=passthrough, key_converters=key_converters)
+        attr = _cache_attr_name(func)
+        get_cache = _cache_getter(spec, attr, cache_factory)
+        descriptor = cachetools.cachedmethod(get_cache, key=spec.make_key, lock=lock, condition=condition, info=info)(
+            func
+        )
+        _mark_descriptor(descriptor, attr, on_copy)
+        return descriptor
+
+    return decorate if func is None else decorate(func)
+
+
+def _cache_getter(
+    spec: ArgSpec, attr: str, cache_factory: Callable[[], MutableMapping]
+) -> Callable[[Any], MutableMapping]:
+    """The function cachetools calls to get an instance's cache: made on first use, stored in ``__dict__[attr]``."""
+
+    def get_cache(obj: Any) -> MutableMapping:
+        try:
+            instance_dict = obj.__dict__
+        except AttributeError:
+            raise TypeError(
+                f"{spec.qualname}: {type(obj).__name__} instances have no __dict__ to hold the cache"
+            ) from None
+        cache = instance_dict.get(attr)
+        if cache is None:
+            check_copy_handling(type(obj))
+            cache = instance_dict.setdefault(attr, cache_factory())  # setdefault: another thread may have made one
+        elif not isinstance(cache, MutableMapping):
+            raise TypeError(f"{spec.qualname}: attribute {attr!r} is already used for something other than a cache")
+        return make_view(cache, spec.qualname) if spec.passthrough else cache
+
+    return get_cache
+
+
+def _mark_descriptor(descriptor: Any, attr: str, on_copy: OnCopy | None) -> None:
+    """Add our checks to cachetools' descriptor, and the markers ``clear_caches`` and copying look for."""
+    add_check(descriptor, _DescriptorChecks)
+    setattr(descriptor, CACHE_ATTR, attr)
+    setattr(descriptor, ON_COPY, ("empty", False) if on_copy is None else (on_copy, True))
+
+
+def _cache_attr_name(func: Callable[..., Any]) -> str:
+    """``_<Class>_<method>_cache`` from the function's qualified name, e.g. ``Sim.draw`` -> ``_Sim_draw_cache``."""
+    qualname = getattr(func, "__qualname__", func.__name__)
+    qualname = qualname.rsplit("<locals>.", 1)[-1]  # drop the enclosing function for local classes
+    return "_" + re.sub(r"\W", "_", qualname.replace(".", "_")) + "_cache"

mcrotk_cache_utils/decorators/rng_caching.py

@@ -0,0 +1,194 @@
+"""Caching decorators for methods and functions that take a random number generator argument.
+
+Two styles of such functions are covered:
+
+- **omit** (``rng_omit_cached_method``, ``rng_omit_cached_function``): the function takes the
+  generator (or a seed) and draws from it, advancing the caller's generator in place, as in the
+  recommended practice (Scientific Python SPEC 7). The generator is just left out of the key.
+- **passthrough** (``rng_passthrough_cached_method``, ``rng_passthrough_cached_function``): the
+  function also returns the generator with its result. On a hit, the caller's generator is
+  returned in place of the one stored with the result.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping
+from typing import TYPE_CHECKING, Any, TypeVar, overload
+
+from ..keys.cache_keys import ArgRef, Passthrough, _as_refs
+from .function_caching import cached_function
+from .method_caching import cached_method
+from .support.decorator_options import (
+    RngOmitFunctionOptions,
+    RngOmitMethodOptions,
+    RngPassthroughFunctionOptions,
+    RngPassthroughMethodOptions,
+)
+
+if TYPE_CHECKING:
+    from typing_extensions import Unpack
+
+F = TypeVar("F", bound=Callable[..., Any])  # the decorated function: its type is kept as it is
+
+
+# --- omit: functions that advance the caller's generator in place ---
+
+
+@overload
+def rng_omit_cached_method(func: F, /, **options: Unpack[RngOmitMethodOptions]) -> F: ...
+@overload
+def rng_omit_cached_method(func: None = None, /, **options: Unpack[RngOmitMethodOptions]) -> Callable[[F], F]: ...
+def rng_omit_cached_method(
+    func: Callable[..., Any] | None = None,
+    /,
+    *,
+    rng_arg: ArgRef = "rng",
+    omit: ArgRef | Iterable[ArgRef] = (),
+    **kwargs: Any,
+) -> Any:
+    """Cache a method that takes a random number generator, leaving the generator out of the key.
+
+    Use as ``@rng_omit_cached_method`` or ``@rng_omit_cached_method(...)``. The same as
+    :func:`cached_method` with ``omit=rng_arg`` (other arguments can be added with ``omit=``);
+    all other options are as for :func:`cached_method`.
+
+    For methods that take the generator, or a seed, as an argument and draw from it::
+
+        @rng_omit_cached_method
+        def simulate(self, n, rng=None):
+            rng = np.random.default_rng(rng)
+            ...
+            return result
+
+    - Calls differing only in their generator share a cached result: the first call's random
+      draws are reused.
+    - On a miss, the method draws from the caller's generator, advancing it in place (when a
+      ``Generator`` is passed). On a hit, nothing is drawn, so the generator is unchanged.
+
+    Args:
+        rng_arg: The generator argument, by name or position (positions count ``self``).
+            Default ``'rng'``.
+    """
+    return cached_method(func, omit=(rng_arg, *_as_refs(omit)), **kwargs)
+
+
+@overload
+def rng_omit_cached_function(func: F, /, **options: Unpack[RngOmitFunctionOptions]) -> F: ...
+@overload
+def rng_omit_cached_function(func: None = None, /, **options: Unpack[RngOmitFunctionOptions]) -> Callable[[F], F]: ...
+def rng_omit_cached_function(
+    func: Callable[..., Any] | None = None,
+    /,
+    *,
+    rng_arg: ArgRef = "rng",
+    omit: ArgRef | Iterable[ArgRef] = (),
+    **kwargs: Any,
+) -> Any:
+    """Cache a function that takes a random number generator, leaving the generator out of the key.
+
+    The function version of :func:`rng_omit_cached_method`, built on :func:`cached_function`
+    (whose options it takes). Also for static and class methods, with ``@staticmethod`` or
+    ``@classmethod`` above it.
+
+    Args:
+        rng_arg: The generator argument, by name or position (as in the ``def`` line, from 0).
+            Default ``'rng'``.
+    """
+    return cached_function(func, omit=(rng_arg, *_as_refs(omit)), **kwargs)
+
+
+# --- passthrough: functions that also return the generator ---
+
+
+@overload
+def rng_passthrough_cached_method(func: F, /, **options: Unpack[RngPassthroughMethodOptions]) -> F: ...
+@overload
+def rng_passthrough_cached_method(
+    func: None = None, /, **options: Unpack[RngPassthroughMethodOptions]
+) -> Callable[[F], F]: ...
+def rng_passthrough_cached_method(
+    func: Callable[..., Any] | None = None,
+    /,
+    *,
+    rng_arg: ArgRef = "rng",
+    rng_output: int = -1,
+    passthrough: Mapping[ArgRef, int | Passthrough] | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Cache a method that takes a random number generator and returns it, leaving it out of the key.
+
+    Use as ``@rng_passthrough_cached_method`` or ``@rng_passthrough_cached_method(...)``. The
+    same as :func:`cached_method` with
+    ``passthrough={rng_arg: (rng_output, numpy.random.default_rng)}`` (other passthrough
+    entries can be added with ``passthrough=``); all other options are as for
+    :func:`cached_method`. Needs numpy.
+
+    For methods that take the generator, or a seed, and also return the generator with their
+    result, so that a caller who passed a seed or ``None`` can continue the stream::
+
+        @rng_passthrough_cached_method
+        def simulate(self, n, rng=None):
+            rng = np.random.default_rng(rng)
+            ...
+            return result, rng
+
+    - Calls differing only in their generator share a cached result: the first call's random
+      draws are reused.
+    - On a miss, the method's own outputs are returned, including the generator it advanced.
+    - On a hit, the stored result is returned with the caller's generator passed through
+      ``numpy.random.default_rng``: a ``Generator`` comes back unchanged (nothing was
+      drawn), and a seed or ``None`` comes back as a new ``Generator`` made from it.
+
+    For methods that don't return the generator, use :func:`rng_omit_cached_method`.
+
+    Args:
+        rng_arg: The generator argument, by name or position (positions count ``self``).
+            Default ``'rng'``.
+        rng_output: Position of the generator in the returned tuple. Default ``-1`` (last).
+    """
+    combined = _rng_passthrough("rng_passthrough_cached_method", rng_arg, rng_output, passthrough)
+    return cached_method(func, passthrough=combined, **kwargs)
+
+
+@overload
+def rng_passthrough_cached_function(func: F, /, **options: Unpack[RngPassthroughFunctionOptions]) -> F: ...
+@overload
+def rng_passthrough_cached_function(
+    func: None = None, /, **options: Unpack[RngPassthroughFunctionOptions]
+) -> Callable[[F], F]: ...
+def rng_passthrough_cached_function(
+    func: Callable[..., Any] | None = None,
+    /,
+    *,
+    rng_arg: ArgRef = "rng",
+    rng_output: int = -1,
+    passthrough: Mapping[ArgRef, int | Passthrough] | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Cache a function that takes a random number generator and returns it, leaving it out of the key.
+
+    The function version of :func:`rng_passthrough_cached_method`, built on
+    :func:`cached_function` (whose options it takes). Also for static and class methods, with
+    ``@staticmethod`` or ``@classmethod`` above it. Needs numpy.
+
+    Args:
+        rng_arg: The generator argument, by name or position (as in the ``def`` line, from 0).
+            Default ``'rng'``.
+        rng_output: Position of the generator in the returned tuple. Default ``-1`` (last).
+    """
+    combined = _rng_passthrough("rng_passthrough_cached_function", rng_arg, rng_output, passthrough)
+    return cached_function(func, passthrough=combined, **kwargs)
+
+
+def _rng_passthrough(
+    decorator: str, rng_arg: ArgRef, rng_output: int, passthrough: Mapping[ArgRef, int | Passthrough] | None
+) -> dict[ArgRef, int | Passthrough]:
+    """The ``passthrough`` option for a passthrough RNG decorator: the RNG entry plus any others given."""
+    try:
+        from numpy.random import default_rng
+    except ImportError as e:
+        raise ImportError(f"{decorator} needs numpy: pip install 'mcrotk-cache-utils[numpy]'") from e
+
+    if passthrough is not None and rng_arg in passthrough:
+        raise ValueError(f"{decorator}: {rng_arg!r} is set by rng_arg/rng_output; don't also put it in passthrough")
+    return {rng_arg: (rng_output, default_rng), **(passthrough or {})}

mcrotk_cache_utils/decorators/support/__init__.py

@@ -0,0 +1 @@
+"""What the decorators are built from: their options (for type checkers) and the placement checks."""

mcrotk_cache_utils/decorators/support/decorator_options.py

@@ -0,0 +1,55 @@
+"""The decorators' options, defined once for type checkers.
+
+The decorators' typed overloads take ``**options: Unpack[...]`` of these, so each option's name
+and type is written in one place. Only type checkers use them; the defaults are in the
+decorators' own signatures.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping, MutableMapping
+from typing import Any, TypedDict
+
+from ...copying.copy_handling import OnCopy
+from ...keys.cache_keys import ArgRef, Passthrough
+
+
+class _Options(TypedDict, total=False):
+    """The options of both ``cached_method`` and ``cached_function``."""
+
+    cache_factory: Callable[[], MutableMapping]
+    omit: ArgRef | Iterable[ArgRef]
+    passthrough: Mapping[ArgRef, int | Passthrough] | None
+    key_converters: Mapping[ArgRef, Callable[[Any], Any]] | None
+    info: bool
+
+
+class MethodOptions(_Options, total=False):
+    """``cached_method``'s options: ``lock`` and ``condition`` are functions of ``self``."""
+
+    lock: Callable[[Any], Any] | None
+    condition: Callable[[Any], Any] | None
+    on_copy: OnCopy | None
+
+
+class FunctionOptions(_Options, total=False):
+    """``cached_function``'s options: ``lock`` and ``condition`` are the objects themselves."""
+
+    lock: Any
+    condition: Any
+
+
+class RngOmitMethodOptions(MethodOptions, total=False):
+    rng_arg: ArgRef
+
+
+class RngOmitFunctionOptions(FunctionOptions, total=False):
+    rng_arg: ArgRef
+
+
+class RngPassthroughMethodOptions(RngOmitMethodOptions, total=False):
+    rng_output: int
+
+
+class RngPassthroughFunctionOptions(RngOmitFunctionOptions, total=False):
+    rng_output: int

mcrotk_cache_utils/decorators/support/placement_checks.py

@@ -0,0 +1,61 @@
+"""Clear errors when a decorator is used in the wrong place (e.g. ``cached_function`` on a method).
+
+cachetools' wrappers and descriptors are ordinary instances, so we give each one a subclass
+of its own class that adds the check, rather than putting another layer in front of it.
+"""
+
+import functools
+from typing import Any
+
+
+def add_check(obj: Any, check: type) -> None:
+    """Make ``obj`` an instance of a subclass of its class with the ``check`` mixin in front."""
+    obj.__class__ = _with_check(type(obj), check)
+
+
+@functools.cache
+def _with_check(cls: type, check: type) -> type:
+    """A subclass of ``cls`` with the ``check`` mixin in front (cached: one class per pair)."""
+    # Same name and module as cachetools' class, so reprs and error messages don't change.
+    return type(cls.__name__, (check, cls), {"__module__": cls.__module__, "__qualname__": cls.__qualname__})
+
+
+class FunctionNotInClassBody:
+    """Mixin for ``cached_function``'s wrapper: an error if it's placed directly in a class body.
+
+    Python calls ``__set_name__`` on objects in a class body when the class is created, but not
+    through ``staticmethod`` or ``classmethod``, so those (the right way to cache static and
+    class methods) are unaffected.
+    """
+
+    def __set_name__(self, owner: type, name: str) -> None:
+        raise TypeError(
+            f"{owner.__qualname__}.{name}: @cached_function can't cache a method: its one cache would be "
+            "shared by all instances, with self in the key. Use @cached_method instead (a cache per instance). "
+            "For a static or class method, put @staticmethod/@classmethod above @cached_function."
+        )
+
+
+class MethodInClassBody:
+    """Mixin for ``cached_method``'s descriptor: an error if it's called directly.
+
+    cachetools' method descriptor is only called directly when it never became a method:
+    under ``@staticmethod`` or ``@classmethod`` (which don't call ``__set_name__``), or on a
+    function outside a class. cachetools' own error there only mentions class methods.
+    """
+
+    def __call__(self, /, *args: Any, **kwargs: Any) -> Any:
+        qualname = getattr(self, "__qualname__", "cached method")
+        raise TypeError(
+            f"{qualname}: @cached_method only works on ordinary methods defined in a class body "
+            "(it keeps a cache per instance). For a function, a static method or a class method, "
+            "use @cached_function, with @staticmethod/@classmethod above it."
+        )
+
+
+def reject_static_or_class_method(func: Any, decorator: str, advice: str) -> None:
+    """Error for ``@decorator`` placed above ``@staticmethod``/``@classmethod`` (it receives their object)."""
+    if isinstance(func, (staticmethod, classmethod)):
+        kind = type(func).__name__
+        qualname = getattr(func, "__qualname__", repr(func))
+        raise TypeError(f"{qualname}: @{decorator} is above @{kind}. {advice}")

mcrotk_cache_utils/keys/__init__.py

@@ -0,0 +1 @@
+"""How a call's arguments become a cache entry: key building, passthrough, and the key converters."""