Propagate attrs with unary, binary functions

Closes #3490
Closes #4065
Closes #3433
Closes #3595

Author: dcherian

xarray/core/arithmetic.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from .options import OPTIONS
+from .options import OPTIONS, _get_keep_attrs
 from .pycompat import dask_array_type
 from .utils import not_implemented
 
@@ -77,6 +77,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             dataset_fill_value=np.nan,
             kwargs=kwargs,
             dask="allowed",
+            keep_attrs=_get_keep_attrs(default=False),
         )
 
     # this has no runtime function - these are listed so IDEs know these

xarray/core/computation.py

@@ -223,8 +223,13 @@ def apply_dataarray_vfunc(
             args, join=join, copy=False, exclude=exclude_dims, raise_on_invalid=False
         )
 
-    if keep_attrs and hasattr(args[0], "name"):
-        name = args[0].name
+    for arg in args:
+        first_obj = arg
+        if isinstance(arg, DataArray):
+            break
+
+    if keep_attrs and hasattr(first_obj, "name"):
+        name = first_obj.name
     else:
         name = result_name(args)
     result_coords = build_output_coords(args, signature, exclude_dims)
@@ -241,6 +246,12 @@ def apply_dataarray_vfunc(
         (coords,) = result_coords
         out = DataArray(result_var, coords, name=name, fastpath=True)
 
+    if keep_attrs and hasattr(first_obj, "attrs"):
+        if isinstance(out, tuple):
+            out = tuple(da._copy_attrs_from(first_obj) for da in out)
+        else:
+            out._copy_attrs_from(first_obj)
+
     return out
 
 
@@ -361,7 +372,10 @@ def apply_dataset_vfunc(
     """
     from .dataset import Dataset
 
-    first_obj = args[0]  # we'll copy attrs from this in case keep_attrs=True
+    for arg in args:
+        first_obj = args
+        if isinstance(first_obj, Dataset):
+            break
 
     if dataset_join not in _JOINS_WITHOUT_FILL_VALUES and fill_value is _NO_FILL_VALUE:
         raise TypeError(
@@ -554,6 +568,11 @@ def apply_variable_ufunc(
     """
     from .variable import Variable, as_compatible_data
 
+    for arg in args:
+        first_obj = arg
+        if isinstance(arg, Variable):
+            break
+
     dim_sizes = unified_dim_sizes(
         (a for a in args if hasattr(a, "dims")), exclude_dims=exclude_dims
     )
@@ -639,8 +658,8 @@ def func(*arrays):
                     )
                 )
 
-        if keep_attrs and isinstance(args[0], Variable):
-            var.attrs.update(args[0].attrs)
+        if keep_attrs and isinstance(first_obj, Variable):
+            var.attrs.update(first_obj.attrs)
         output.append(var)
 
     if signature.num_outputs == 1:

xarray/core/dataset.py

@@ -4368,12 +4368,15 @@ def map(
             foo      (dim_0, dim_1) float64 0.3751 1.951 1.945 0.2948 0.711 0.3948
             bar      (x) float64 1.0 2.0
         """
+        if keep_attrs is None:
+            keep_attrs = _get_keep_attrs(default=False)
         variables = {
             k: maybe_wrap_array(v, func(v, *args, **kwargs))
             for k, v in self.data_vars.items()
         }
-        if keep_attrs is None:
-            keep_attrs = _get_keep_attrs(default=False)
+        if keep_attrs:
+            for k, v in variables.items():
+                v._copy_attrs_from(self.data_vars[k])
         attrs = self.attrs if keep_attrs else None
         return type(self)(variables, attrs=attrs)
 

xarray/core/variable.py

@@ -2024,7 +2024,9 @@ def imag(self):
         return type(self)(self.dims, self.data.imag, self._attrs)
 
     def __array_wrap__(self, obj, context=None):
-        return Variable(self.dims, obj)
+        keep_attrs = _get_keep_attrs(default=False)
+        attrs = self._attrs if keep_attrs else {}
+        return Variable(self.dims, obj, attrs)
 
     @staticmethod
     def _unary_op(f):

xarray/tests/test_dataarray.py

@@ -9,7 +9,15 @@
 import pytest
 
 import xarray as xr
-from xarray import DataArray, Dataset, IndexVariable, Variable, align, broadcast
+from xarray import (
+    DataArray,
+    Dataset,
+    IndexVariable,
+    Variable,
+    align,
+    broadcast,
+    set_options,
+)
 from xarray.coding.times import CFDatetimeCoder
 from xarray.convert import from_cdms2
 from xarray.core import dtypes
@@ -2458,6 +2466,21 @@ def test_assign_attrs(self):
         assert_identical(new_actual, expected)
         assert actual.attrs == {"a": 1, "b": 2}
 
+    def test_propagate_attrs(self):
+        da = DataArray(self.va)
+
+        # test defaults
+        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
+
+        with set_options(keep_attrs=True):
+            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
+
     def test_fillna(self):
         a = DataArray([np.nan, 1, np.nan, 3], coords={"x": range(4)}, dims="x")
         actual = a.fillna(-1)