Propagate attrs with DataArray unary ops

Author: dcherian

xarray/core/dataarray.py

@@ -55,7 +55,7 @@
 from .indexes import Indexes, default_indexes, propagate_indexes
 from .indexing import is_fancy_indexer
 from .merge import PANDAS_TYPES, MergeError, _extract_indexes_from_coords
-from .options import OPTIONS
+from .options import OPTIONS, _get_keep_attrs
 from .utils import Default, ReprObject, _check_inplace, _default, either_dict_or_kwargs
 from .variable import (
     IndexVariable,
@@ -2730,13 +2730,19 @@ def __rmatmul__(self, other):
     def _unary_op(f: Callable[..., Any]) -> Callable[..., "DataArray"]:
         @functools.wraps(f)
         def func(self, *args, **kwargs):
+            keep_attrs = kwargs.pop("keep_attrs", None)
+            if keep_attrs is None:
+                keep_attrs = _get_keep_attrs(default=True)
             with warnings.catch_warnings():
                 warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")
                 warnings.filterwarnings(
                     "ignore", r"Mean of empty slice", category=RuntimeWarning
                 )
                 with np.errstate(all="ignore"):
-                    return self.__array_wrap__(f(self.variable.data, *args, **kwargs))
+                    da = self.__array_wrap__(f(self.variable.data, *args, **kwargs))
+                if keep_attrs:
+                    da.attrs = self.attrs
+                return da
 
         return func
 

xarray/tests/test_dataarray.py

@@ -2488,21 +2488,20 @@ def test_assign_attrs(self):
         assert_identical(new_actual, expected)
         assert actual.attrs == {"a": 1, "b": 2}
 
-    def test_propagate_attrs(self):
+    @pytest.mark.parametrize(
+        "func", [lambda x: x.clip(0, 1), lambda x: np.float64(1.0) * x, np.abs, abs]
+    )
+    def test_propagate_attrs(self, func):
         da = DataArray(self.va)
 
         # test defaults
-        assert (np.float64(1.0) * da).attrs == da.attrs
-        assert np.abs(da).attrs == da.attrs
-        # not sure about the next two
-        assert da.clip(0, 1).attrs != da.attrs
-        assert abs(da).attrs != da.attrs
+        assert func(da).attrs == da.attrs
 
         with set_options(keep_attrs=False):
-            assert da.clip(0, 1).attrs != da.attrs
-            assert (np.float64(1.0) * da).attrs != da.attrs
-            assert np.abs(da).attrs != da.attrs
-            assert abs(da).attrs != da.attrs
+            assert func(da).attrs == {}
+
+        with set_options(keep_attrs=True):
+            assert func(da).attrs == da.attrs
 
     def test_fillna(self):
         a = DataArray([np.nan, 1, np.nan, 3], coords={"x": range(4)}, dims="x")