Add a context manager as an alternative to monkey patching (#91)

* Add a context manager for default dims

* Add a docstring and fix lint issues

* Add default linalg dims to package exports

* add type hints

* Handle all iterables

* Add missing contextmanager wrapper

* Resolve PR comments

Author: krokosik

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)

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]: ...

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,