diff --git a/src/xarray_einstats/linalg.py b/src/xarray_einstats/linalg.py index 0bead5b..aef88ce 100644 --- a/src/xarray_einstats/linalg.py +++ b/src/xarray_einstats/linalg.py @@ -8,10 +8,14 @@ example usage. The functions that are not available via the accessor are ``einsum``, ``einsum_path``, - ``matmul`` and ``get_default_dims``. + ``matmul``, ``get_default_dims`` and ``default_dims``. """ +import sys +from collections.abc import Iterable +from contextlib import contextmanager + import numpy as np import xarray as xr @@ -35,6 +39,7 @@ "solve", "inv", "pinv", + "default_dims", ] @@ -109,6 +114,10 @@ def get_default_dims(dims1, dims2): You can still use ``dims`` explicitly to override those defaults. + .. note:: + Monkeypatching ``get_default_dims`` directly works but is error-prone. + Consider using the :func:`default_dims` context manager instead. + """ raise MissingMonkeypatchError() @@ -119,12 +128,60 @@ def _attempt_default_dims(func, da1_dims, da2_dims=None): aux = get_default_dims(da1_dims, da2_dims) except MissingMonkeypatchError: raise TypeError( - f"{func} missing required argument dims. You must monkeypatch " - "xarray_einstats.linalg.get_default_dims for dims=None to be supported" + f"{func} missing required argument dims. Use " + "xarray_einstats.linalg.default_dims context manager or pass dims explicitly" ) from None return aux +@contextmanager +def default_dims(func_or_dims): + """Context manager to temporarily set the default dimensions for linalg functions. + + Safer alternative to monkey patching :func:`get_default_dims`, + as it ensures that the original function is restored even if an error occurs + within the context. + + Parameters + ---------- + func_or_dims : callable or iterable + If a callable is provided, it should take the same arguments as :func:`get_default_dims` + and return the default dimensions based on those arguments. + If an iterable is provided, it will be used as the default dimensions + regardless of the input arguments. + + See Also + -------- + get_default_dims + + Examples + -------- + Set the default dims to ``("dim", "dim2")`` for the duration of the ``with`` block: + + .. code-block:: python + + from xarray_einstats import linalg, tutorial + da = tutorial.generate_matrices_dataarray(5) + + with linalg.default_dims(("dim", "dim2")): + linalg.inv(da) + + """ + _linalg = sys.modules[__name__] + original_get_default_dims = _linalg.get_default_dims + + def func(*args): + if isinstance(func_or_dims, Iterable): + return func_or_dims + return func_or_dims(*args) + + _linalg.get_default_dims = func + try: + yield + finally: + _linalg.get_default_dims = original_get_default_dims + + class PairHandler: def __init__(self, all_dims, keep_dims): self.potential_out_dims = keep_dims.union(all_dims) diff --git a/src/xarray_einstats/linalg.pyi b/src/xarray_einstats/linalg.pyi index d7ef47c..45c2771 100644 --- a/src/xarray_einstats/linalg.pyi +++ b/src/xarray_einstats/linalg.pyi @@ -1,12 +1,12 @@ # File generated with docstub import numbers -from collections.abc import Hashable, Sequence -from typing import Literal +from collections.abc import Hashable, Iterable, Sequence +from contextlib import contextmanager +from typing import Callable, Generator, Literal import numpy as np import xarray -import xarray as xr from _typeshed import Incomplete from numpy.typing import NDArray @@ -30,6 +30,7 @@ __all__ = [ "solve", "inv", "pinv", + "default_dims", ] class MissingMonkeypatchError(Exception): @@ -195,3 +196,7 @@ def pinv( hermitian: bool = ..., **kwargs: Incomplete, ) -> xarray.DataArray: ... +@contextmanager +def default_dims( + func_or_dims: Callable | Iterable, +) -> Generator[None, None, None]: ... diff --git a/tests/test_linalg.py b/tests/test_linalg.py index f688f5b..d0e00dc 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -79,6 +79,19 @@ def default_dims(dims1, dims2): # pylint: disable=unused-argument assert out.dims == matrices.dims +def test_default_dims_context_manager(matrices): + with pytest.raises(TypeError, match="missing required argument dims"): + inv(matrices) + + with linalg.default_dims(("dim", "dim2")): + out = inv(matrices) + assert out.dims == matrices.dims + + # outside the context, it should raise again + with pytest.raises(TypeError, match="missing required argument dims"): + inv(matrices) + + class TestEinsumFamily: # raw_einsum calls einsum, so the tests on raw_einsum also cover einsum, then # there are some specific ones for various reasons,