apply_ufunc: don't modify attrs on input variables (pydata#10330)

Co-authored-by: Erik Månsson <none@example>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Author: 4 people

xarray/structure/merge.py

@@ -283,11 +283,17 @@ def merge_collected(
                                 "conflicting attribute values on combined "
                                 f"variable {name!r}:\nfirst value: {variable.attrs!r}\nsecond value: {other_variable.attrs!r}"
                             )
-                merged_vars[name] = variable
-                merged_vars[name].attrs = merge_attrs(
+                attrs = merge_attrs(
                     [var.attrs for var, _ in indexed_elements],
                     combine_attrs=combine_attrs,
                 )
+                if variable.attrs or attrs:
+                    # Make a shallow copy to so that assigning merged_vars[name].attrs
+                    # does not affect the original input variable.
+                    merged_vars[name] = variable.copy(deep=False)
+                    merged_vars[name].attrs = attrs
+                else:
+                    merged_vars[name] = variable
                 merged_indexes[name] = index
             else:
                 variables = [variable for variable, _ in elements_list]

xarray/tests/test_computation.py

@@ -2206,6 +2206,7 @@ def test_where() -> None:
 def test_where_attrs() -> None:
     cond = xr.DataArray([True, False], coords={"a": [0, 1]}, attrs={"attr": "cond_da"})
     cond["a"].attrs = {"attr": "cond_coord"}
+    input_cond = cond.copy()
     x = xr.DataArray([1, 1], coords={"a": [0, 1]}, attrs={"attr": "x_da"})
     x["a"].attrs = {"attr": "x_coord"}
     y = xr.DataArray([0, 0], coords={"a": [0, 1]}, attrs={"attr": "y_da"})
@@ -2216,6 +2217,22 @@ def test_where_attrs() -> None:
     expected = xr.DataArray([1, 0], coords={"a": [0, 1]}, attrs={"attr": "x_da"})
     expected["a"].attrs = {"attr": "x_coord"}
     assert_identical(expected, actual)
+    # Check also that input coordinate attributes weren't modified by reference
+    assert x["a"].attrs == {"attr": "x_coord"}
+    assert y["a"].attrs == {"attr": "y_coord"}
+    assert cond["a"].attrs == {"attr": "cond_coord"}
+    assert_identical(cond, input_cond)
+
+    # 3 DataArrays, drop attrs
+    actual = xr.where(cond, x, y, keep_attrs=False)
+    expected = xr.DataArray([1, 0], coords={"a": [0, 1]})
+    assert_identical(expected, actual)
+    assert_identical(expected.coords["a"], actual.coords["a"])
+    # Check also that input coordinate attributes weren't modified by reference
+    assert x["a"].attrs == {"attr": "x_coord"}
+    assert y["a"].attrs == {"attr": "y_coord"}
+    assert cond["a"].attrs == {"attr": "cond_coord"}
+    assert_identical(cond, input_cond)
 
     # x as a scalar, takes no attrs
     actual = xr.where(cond, 0, y, keep_attrs=True)

xarray/tests/test_merge.py

@@ -183,9 +183,11 @@ def test_merge_arrays_attrs_variables(
         self, combine_attrs, attrs1, attrs2, expected_attrs, expect_exception
     ):
         """check that combine_attrs is used on data variables and coords"""
+        input_attrs1 = attrs1.copy()
         data1 = xr.Dataset(
             {"var1": ("dim1", [], attrs1)}, coords={"dim1": ("dim1", [], attrs1)}
         )
+        input_attrs2 = attrs2.copy()
         data2 = xr.Dataset(
             {"var1": ("dim1", [], attrs2)}, coords={"dim1": ("dim1", [], attrs2)}
         )
@@ -202,6 +204,12 @@ def test_merge_arrays_attrs_variables(
 
             assert_identical(actual, expected)
 
+            # Check also that input attributes weren't modified
+            assert data1["var1"].attrs == input_attrs1
+            assert data1.coords["dim1"].attrs == input_attrs1
+            assert data2["var1"].attrs == input_attrs2
+            assert data2.coords["dim1"].attrs == input_attrs2
+
     def test_merge_attrs_override_copy(self):
         ds1 = xr.Dataset(attrs={"x": 0})
         ds2 = xr.Dataset(attrs={"x": 1})
@@ -344,6 +352,18 @@ def test_merge(self):
         with pytest.raises(ValueError, match=r"should be coordinates or not"):
             data.merge(data.reset_coords())
 
+    def test_merge_drop_attrs(self):
+        data = create_test_data()
+        ds1 = data[["var1"]]
+        ds2 = data[["var3"]]
+        ds1.coords["dim2"].attrs["keep me"] = "example"
+        ds2.coords["numbers"].attrs["foo"] = "bar"
+        actual = ds1.merge(ds2, combine_attrs="drop")
+        assert actual.coords["dim2"].attrs == {}
+        assert actual.coords["numbers"].attrs == {}
+        assert ds1.coords["dim2"].attrs["keep me"] == "example"
+        assert ds2.coords["numbers"].attrs["foo"] == "bar"
+
     def test_merge_broadcast_equals(self):
         ds1 = xr.Dataset({"x": 0})
         ds2 = xr.Dataset({"x": ("y", [0, 0])})

xarray/tests/test_ufuncs.py

@@ -62,6 +62,19 @@ def test_binary_out():
         assert_identical(actual_exponent, arg)
 
 
+def test_binary_coord_attrs():
+    t = xr.Variable("t", np.arange(2, 4), attrs={"units": "s"})
+    x = xr.DataArray(t.values**2, coords={"t": t}, attrs={"units": "s^2"})
+    y = xr.DataArray(t.values**3, coords={"t": t}, attrs={"units": "s^3"})
+    z1 = xr.apply_ufunc(np.add, x, y, keep_attrs=True)
+    assert z1.coords["t"].attrs == {"units": "s"}
+    z2 = xr.apply_ufunc(np.add, x, y, keep_attrs=False)
+    assert z2.coords["t"].attrs == {}
+    # Check also that input array's coordinate attributes weren't affected
+    assert t.attrs == {"units": "s"}
+    assert x.coords["t"].attrs == {"units": "s"}
+
+
 def test_groupby():
     ds = xr.Dataset({"a": ("x", [0, 0, 0])}, {"c": ("x", [0, 0, 1])})
     ds_grouped = ds.groupby("c")