ENH: improve hash behavior with cache_to_disk (#157)

* BREAK: make `backend` argument keyword only
* ENH: embed `cloudpickle` and `ampform` as `lambdify` dependencies
* ENH: embed jax and sympy version in cache dependency
* MAINT: improve docstring positioning

Author: redeboer

src/ampform_dpd/io/cached.py

@@ -2,17 +2,10 @@
 
 from __future__ import annotations
 
-import logging
-import pickle
-from importlib.metadata import version
-from pathlib import Path
 from typing import TYPE_CHECKING, overload
 
 import cloudpickle
-from ampform.sympy._cache import (
-    get_readable_hash,  # noqa: PLC2701
-    get_system_cache_directory,  # noqa: PLC2701
-)
+from ampform.sympy._cache import cache_to_disk  # noqa: PLC2701
 from ampform.sympy.cached import (
     doit,  # noqa: F401  # pyright: ignore[reportUnusedImport]
     unfold,  # noqa: F401  # pyright: ignore[reportUnusedImport]
@@ -22,7 +15,6 @@
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
-    from typing import Any
 
     import sympy as sp
     from tensorwaves.function import (
@@ -31,27 +23,25 @@
     )
     from tensorwaves.interface import Function, ParameterValue, ParametrizedFunction
 
-_LOGGER = logging.getLogger(__name__)
-
 
 @overload
-def lambdify(
-    expr: sp.Expr,
-    backend: str = "jax",
-    directory: str | None = None,
-) -> PositionalArgumentFunction: ...
+def lambdify(expr: sp.Expr, *, backend: str = "jax") -> PositionalArgumentFunction: ...
 @overload
 def lambdify(
     expr: sp.Expr,
     parameters: Mapping[sp.Symbol, ParameterValue],
+    *,
     backend: str = "jax",
-    directory: str | None = None,
 ) -> ParametrizedBackendFunction: ...
-def lambdify(  # type:ignore[misc]  # pyright:ignore[reportInconsistentOverload]
+@cache_to_disk(
+    dump_function=cloudpickle.dump,
+    dependencies=["cloudpickle", "ampform", "jax", "sympy"],
+)
+def lambdify(
     expr: sp.Expr,
     parameters: Mapping[sp.Symbol, ParameterValue] | None = None,
+    *,
     backend: str = "jax",
-    cache_directory: Path | str | None = None,
 ) -> ParametrizedFunction | Function:
     """Lambdify a SymPy `~sympy.core.expr.Expr` and cache the result to disk.
 
@@ -65,41 +55,11 @@ def lambdify(  # type:ignore[misc]  # pyright:ignore[reportInconsistentOverload]
         parameters: Specify this argument in order to create a
             `~tensorwaves.function.ParametrizedBackendFunction` instead of a
             `~tensorwaves.function.PositionalArgumentFunction`.
-        backend: The choice of backend for the created numerical function. **WARNING**:
-            this function has only been tested for :code:`backend="jax"`!
-        directory: The directory in which to cache the result. If `None`, the cache
-            directory will be put under the home directory, or to the path specified by
-            the environment variable :code:`SYMPY_CACHE_DIR`.
+        backend: The choice of backend for the created numerical function.
+            **WARNING**: this function has only been tested for :code:`backend="jax"`!
 
     .. seealso:: :func:`ampform.sympy.perform_cached_doit`
     """
-    if cache_directory is None:
-        system_cache_dir = get_system_cache_directory()
-        backend_version = version(backend)
-        cache_directory = (
-            Path(system_cache_dir) / "ampform_dpd" / f"{backend}-v{backend_version}"
-        )
-    if not isinstance(cache_directory, Path):
-        cache_directory = Path(cache_directory)
-    cache_directory.mkdir(exist_ok=True, parents=True)
-    if parameters is None:
-        hash_obj: Any = expr
-    else:
-        hash_obj = (
-            expr,
-            tuple((s, parameters[s]) for s in sorted(parameters, key=str)),
-        )
-    h = get_readable_hash(hash_obj)
-    filename = cache_directory / f"{h}.pkl"
-    if filename.exists():
-        with open(filename, "rb") as f:
-            return pickle.load(f)
-    _LOGGER.warning(f"Cached function file {filename} not found, lambdifying...")
-    func: ParametrizedFunction | Function
     if parameters is None:
-        func = create_function(expr, backend)
-    else:
-        func = create_parametrized_function(expr, parameters, backend)
-    with open(filename, "wb") as f:
-        cloudpickle.dump(func, f)
-    return func
+        return create_function(expr, backend)
+    return create_parametrized_function(expr, parameters, backend)