convert apply_to_dataset to a top-level function

Author: keewis

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/common.py

@@ -635,21 +635,6 @@ def pipe(
         else:
             return func(self, *args, **kwargs)
 
-    def apply_to_dataset(self, f):
-        from .dataarray import DataArray
-
-        if isinstance(self, DataArray):
-            ds = self._to_temp_dataset()
-        else:
-            ds = self
-
-        result = f(ds)
-
-        if isinstance(self, DataArray):
-            return self._from_temp_dataset(result, name=self.name)
-        else:
-            return result
-
     def groupby(self, group, squeeze: bool = True, restore_coord_dims: bool = None):
         """Returns a GroupBy object for performing grouped operations.
 

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}

xarray/tests/test_dataarray.py

@@ -2617,17 +2617,6 @@ def test_fillna(self):
         actual = a.groupby("b").fillna(DataArray([0, 2], dims="b"))
         assert_identical(expected, actual)
 
-    def test_apply_to_dataset(self):
-        def func(ds):
-            return Dataset(ds.data_vars, coords=ds.coords)
-
-        da = DataArray(
-            [[0, 1], [2, 3], [4, 5]],
-            dims=("x", "y"),
-            name="abc",
-        )
-        assert_identical(da, da.apply_to_dataset(func))
-
     def test_groupby_iter(self):
         for ((act_x, act_dv), (exp_x, exp_ds)) in zip(
             self.dv.groupby("y"), self.ds.groupby("y")

xarray/tests/test_dataset.py

@@ -5028,13 +5028,6 @@ def test_count(self):
         actual = ds.count()
         assert_identical(expected, actual)
 
-    def test_apply_to_dataset(self):
-        def func(ds):
-            return Dataset(ds.data_vars, coords=ds.coords)
-
-        ds = create_test_data()
-        assert_identical(ds, ds.apply_to_dataset(func))
-
     def test_map(self):
         data = create_test_data()
         data.attrs["foo"] = "bar"