propagate attrs on coords in Dataset.map (pydata#10602)

* check that weighted ops propagate attrs on coords

* propagate attrs on coords in `map` if keep_attrs

* directly check that `map` propagates attrs on coords

* whats-new

Author: keewis

doc/whats-new.rst

@@ -48,6 +48,8 @@ Bug fixes
 - Fix error when encoding an empty :py:class:`numpy.datetime64` array
   (:issue:`10722`, :pull:`10723`). By `Spencer Clark
   <https://github.com/spencerkclark>`_.
+- Propagation coordinate attrs in :py:meth:`xarray.Dataset.map` (:issue:`9317`, :pull:`10602`).
+  By `Justus Magin <https://github.com/keewis>`_.
 
 Documentation
 ~~~~~~~~~~~~~

xarray/core/dataset.py

@@ -6929,11 +6929,22 @@ def map(
             k: maybe_wrap_array(v, func(v, *args, **kwargs))
             for k, v in self.data_vars.items()
         }
+        coord_vars, indexes = merge_coordinates_without_align(
+            [v.coords for v in variables.values()]
+        )
+        coords = Coordinates._construct_direct(coords=coord_vars, indexes=indexes)
+
         if keep_attrs:
             for k, v in variables.items():
                 v._copy_attrs_from(self.data_vars[k])
+
+            for k, v in coords.items():
+                if k not in self.coords:
+                    continue
+                v._copy_attrs_from(self.coords[k])
+
         attrs = self.attrs if keep_attrs else None
-        return type(self)(variables, attrs=attrs)
+        return type(self)(variables, coords=coords, attrs=attrs)
 
     def apply(
         self,

xarray/tests/test_dataset.py

@@ -6236,6 +6236,38 @@ def scale(x, multiple=1):
         expected = data.drop_vars("time")  # time is not used on a data var
         assert_equal(expected, actual)
 
+    def test_map_coords_attrs(self) -> None:
+        ds = xr.Dataset(
+            {
+                "a": (
+                    ["x", "y", "z"],
+                    np.arange(24).reshape(3, 4, 2),
+                    {"attr1": "value1"},
+                ),
+                "b": ("y", np.arange(4), {"attr2": "value2"}),
+            },
+            coords={
+                "x": ("x", np.array([-1, 0, 1]), {"attr3": "value3"}),
+                "z": ("z", list("ab"), {"attr4": "value4"}),
+            },
+        )
+
+        def func(arr):
+            if "y" not in arr.dims:
+                return arr
+
+            # drop attrs from coords
+            return arr.mean(dim="y").drop_attrs()
+
+        expected = ds.mean(dim="y", keep_attrs=True)
+        actual = ds.map(func, keep_attrs=True)
+
+        assert_identical(actual, expected)
+        assert actual["x"].attrs
+
+        ds["x"].attrs["y"] = "x"
+        assert ds["x"].attrs != actual["x"].attrs
+
     def test_apply_pending_deprecated_map(self) -> None:
         data = create_test_data()
         data.attrs["foo"] = "bar"

xarray/tests/test_weighted.py

@@ -770,6 +770,17 @@ def test_weighted_operations_keep_attr_da_in_ds(operation):
     assert data.a.attrs == result.a.attrs
 
 
+def test_weighted_mean_keep_attrs_ds():
+    weights = DataArray(np.random.randn(2))
+    data = Dataset(
+        {"a": (["dim_0", "dim_1"], np.random.randn(2, 2), dict(attr="data"))},
+        coords={"dim_1": ("dim_1", ["a", "b"], {"attr1": "value1"})},
+    )
+
+    result = data.weighted(weights).mean(dim="dim_0", keep_attrs=True)
+    assert data.coords["dim_1"].attrs == result.coords["dim_1"].attrs
+
+
 @pytest.mark.parametrize("operation", ("sum_of_weights", "sum", "mean", "quantile"))
 @pytest.mark.parametrize("as_dataset", (True, False))
 def test_weighted_bad_dim(operation, as_dataset):