Fix dataset attrs

Author: dcherian

xarray/core/dataset.py

@@ -4899,18 +4899,20 @@ def from_dict(cls, d):
         return obj
 
     @staticmethod
-    def _unary_op(f, keep_attrs=None):
-        if keep_attrs is None:
-            keep_attrs = _get_keep_attrs(default=True)
-
+    def _unary_op(f):
         @functools.wraps(f)
         def func(self, *args, **kwargs):
             variables = {}
+            keep_attrs = kwargs.pop("keep_attrs", None)
+            if keep_attrs is None:
+                keep_attrs = _get_keep_attrs(default=True)
             for k, v in self._variables.items():
                 if k in self._coord_names:
                     variables[k] = v
                 else:
                     variables[k] = f(v, *args, **kwargs)
+                    if keep_attrs:
+                        variables[k].attrs = v._attrs
             attrs = self._attrs if keep_attrs else None
             return self._replace_with_new_dims(variables, attrs=attrs)
 

xarray/core/options.py

@@ -71,7 +71,7 @@ def _get_keep_attrs(default):
         return global_choice
     else:
         raise ValueError(
-            "The global option keep_attrs must be one of" " True, False or 'default'."
+            "The global option keep_attrs must be one of True, False or 'default'."
         )
 
 

xarray/core/variable.py

@@ -2035,16 +2035,20 @@ def imag(self):
         return type(self)(self.dims, self.data.imag, self._attrs)
 
     def __array_wrap__(self, obj, context=None):
-        keep_attrs = _get_keep_attrs(default=False)
-        attrs = self._attrs if keep_attrs else {}
-        return Variable(self.dims, obj, attrs)
+        return Variable(self.dims, obj)
 
     @staticmethod
     def _unary_op(f):
         @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 np.errstate(all="ignore"):
-                return self.__array_wrap__(f(self.data, *args, **kwargs))
+                result = self.__array_wrap__(f(self.data, *args, **kwargs))
+                if keep_attrs:
+                    result.attrs = self._attrs
+                return result
 
         return func
 

xarray/tests/test_dataset.py

@@ -4415,22 +4415,19 @@ def test_fillna(self):
         assert actual.a.name == "a"
         assert actual.a.attrs == ds.a.attrs
 
-    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(range(5), name="a", attrs={"attr": "da"})
         ds = Dataset({"a": da}, attrs={"attr": "ds"})
 
         # test defaults
-        assert ds.clip(0, 1).attrs == ds.attrs
-        assert (np.float64(1.0) * ds).attrs == ds.attrs
-        assert np.abs(ds).attrs == ds.attrs
-        assert abs(ds).attrs == ds.attrs
-
+        assert func(ds).attrs == ds.attrs
         with set_options(keep_attrs=False):
-            assert ds.clip(0, 1).attrs != ds.attrs
-            assert (np.float64(1.0) * ds).attrs != ds.attrs
-            assert np.abs(ds).attrs != ds.attrs
-            assert abs(ds).attrs != ds.attrs
+            assert func(ds).attrs != ds.attrs
+            assert func(ds).a.attrs != ds.a.attrs
 
     def test_where(self):
         ds = Dataset({"a": ("x", range(5))})