apply to dataset

xarray/__init__.py

@@ -18,7 +18,15 @@
 from .core.alignment import align, broadcast
 from .core.combine import combine_by_coords, combine_nested
 from .core.common import ALL_DIMS, full_like, ones_like, zeros_like
-from .core.computation import apply_ufunc, corr, cov, dot, polyval, where
+from .core.computation import (
+    apply_to_dataset,
+    apply_ufunc,
+    corr,
+    cov,
+    dot,
+    polyval,
+    where,
+)
 from .core.concat import concat
 from .core.dataarray import DataArray
 from .core.dataset import Dataset
@@ -46,6 +54,7 @@
     # Top-level functions
     "align",
     "apply_ufunc",
+    "apply_to_dataset",
     "as_variable",
     "broadcast",
     "cftime_range",

xarray/core/computation.py

@@ -1142,6 +1142,42 @@ def earth_mover_distance(first_samples,
         return apply_array_ufunc(func, *args, dask=dask)
 
 
+def apply_to_dataset(func, obj, *args, **kwargs):
+    """apply a function expecting a Dataset to a xarray object
+
+    Parameters
+    ----------
+    func : callable
+        A function expecting a Dataset as its first parameter.
+    obj : DataArray or Dataset
+        The dataset to apply ``func`` to. If a ``DataArray``, convert it to a single
+        variable ``Dataset`` first.
+    *args, **kwargs
+        Additional arguments to ``func``
+
+    Returns
+    -------
+    DataArray or Dataset
+        The result of ``func(obj, *args, **kwargs)`` with the same type as ``obj``.
+
+    Notes
+    -----
+    If a ``DataArray``, result will have the same name as ``obj`` but the single data
+    variable in the temporary ``Dataset`` will always have a generic name.
+    """
+    from .dataarray import DataArray
+
+    ds = obj._to_temp_dataset() if isinstance(obj, DataArray) else obj
+
+    result = func(ds, *args, **kwargs)
+
+    return (
+        obj._from_temp_dataset(result, name=obj.name)
+        if isinstance(obj, DataArray)
+        else result
+    )
+
+
 def cov(da_a, da_b, dim=None, ddof=1):
     """
     Compute covariance between two DataArray objects along a shared dimension.

xarray/tests/test_computation.py

@@ -468,6 +468,49 @@ def test_apply_groupby_add():
         add(data_array.groupby("y"), data_array.groupby("x"))
 
 
+@pytest.mark.parametrize(
+    ["obj", "expected"],
+    (
+        pytest.param(
+            xr.DataArray(
+                [0, 1],
+                coords={
+                    "x": ("x", [-1, 1], {"a": 1, "b": 2}),
+                    "u": ("x", [2, 3], {"c": 3}),
+                },
+                dims="x",
+                attrs={"d": 4, "e": 5},
+            ),
+            xr.DataArray([0, 1], coords={"x": [-1, 1], "u": ("x", [2, 3])}, dims="x"),
+            id="DataArray",
+        ),
+        pytest.param(
+            xr.Dataset(
+                {"a": ("x", [1, 2], {"a": 1, "b": 2}), "b": ("x", [0, 1], {"c": 3})},
+                coords={
+                    "x": ("x", [-1, 1], {"d": 4, "e": 5}),
+                    "u": ("x", [2, 3], {"f": 6}),
+                },
+            ),
+            xr.Dataset(
+                {"a": ("x", [1, 2]), "b": ("x", [0, 1])},
+                coords={"x": [-1, 1], "u": ("x", [2, 3])},
+            ),
+            id="Dataset",
+        ),
+    ),
+)
+def test_apply_to_dataset(obj, expected):
+    def clear_all_attrs(ds):
+        new_ds = ds.copy()
+        for var in new_ds.variables.values():
+            var.attrs.clear()
+        new_ds.attrs.clear()
+        return new_ds
+
+    assert_identical(expected, xr.apply_to_dataset(clear_all_attrs, obj))
+
+
 def test_unified_dim_sizes():
     assert unified_dim_sizes([xr.Variable((), 0)]) == {}
     assert unified_dim_sizes([xr.Variable("x", [1]), xr.Variable("x", [1])]) == {"x": 1}