Repository: pydata/xarray

Issue:
keep_attrs=True does not work `apply_ufunc` with xr.Variable
#### Code Sample, a copy-pastable example if possible

`keep_attrs=True` works nice for xr.DataArray, but it is neglected for `xr.Variable`

```python
In [2]: import numpy as np
In [3]: import xarray as xr
In [4]: da = xr.DataArray([0, 1, 2], dims='x', attrs={'foo': 'var'})
In [5]: func = lambda x: x*2
In [6]: xr.apply_ufunc(func, da, keep_attrs=True, input_core_dims=[['x']], outpu
   ...: t_core_dims=[['z']])
Out[6]: 
<xarray.DataArray (z: 3)>  # attrs are tracked for xr.DataArray
array([0, 2, 4])
Dimensions without coordinates: z
Attributes:
    foo:      var

In [7]: xr.apply_ufunc(func, da.variable, keep_attrs=True, input_core_dims=[['x'
   ...: ]], output_core_dims=[['z']])
Out[7]: 
<xarray.Variable (z: 3)>  # attrs are dropped
array([0, 2, 4])
```
#### Problem description
`keep_attrs=True` should work also with `xr.Variable`

#### Expected Output
```python
<xarray.Variable (z: 3)>  # attrs are dropped
array([0, 2, 4])
Attributes:
    foo:      var
```

Discussion:

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 3c2e143b..b177fc70 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -49,6 +49,9 @@ Enhancements
 Bug fixes
 ~~~~~~~~~
 
+- Fixed a bug where `keep_attrs=True` flag was neglected if
+  :py:func:`apply_func` was used with :py:class:`Variable`. (:issue:`2114`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - When assigning a :py:class:`DataArray` to :py:class:`Dataset`, any conflicted
   non-dimensional coordinates of the DataArray are now dropped.
   (:issue:`2068`)
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index f06e90b5..77a52ac0 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -195,7 +195,6 @@ def apply_dataarray_ufunc(func, *args, **kwargs):
     signature = kwargs.pop('signature')
     join = kwargs.pop('join', 'inner')
     exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET)
-    keep_attrs = kwargs.pop('keep_attrs', False)
     if kwargs:
         raise TypeError('apply_dataarray_ufunc() got unexpected keyword '
                         'arguments: %s' % list(kwargs))
@@ -217,11 +216,6 @@ def apply_dataarray_ufunc(func, *args, **kwargs):
         coords, = result_coords
         out = DataArray(result_var, coords, name=name, fastpath=True)
 
-    if keep_attrs and isinstance(args[0], DataArray):
-        if isinstance(out, tuple):
-            out = tuple(ds._copy_attrs_from(args[0]) for ds in out)
-        else:
-            out._copy_attrs_from(args[0])
     return out
 
 
@@ -526,6 +520,7 @@ def apply_variable_ufunc(func, *args, **kwargs):
     dask = kwargs.pop('dask', 'forbidden')
     output_dtypes = kwargs.pop('output_dtypes', None)
     output_sizes = kwargs.pop('output_sizes', None)
+    keep_attrs = kwargs.pop('keep_attrs', False)
     if kwargs:
         raise TypeError('apply_variable_ufunc() got unexpected keyword '
                         'arguments: %s' % list(kwargs))
@@ -567,11 +562,17 @@ def apply_variable_ufunc(func, *args, **kwargs):
     if signature.num_outputs > 1:
         output = []
         for dims, data in zip(output_dims, result_data):
-            output.append(Variable(dims, data))
+            var = Variable(dims, data)
+            if keep_attrs and isinstance(args[0], Variable):
+                var.attrs.update(args[0].attrs)
+            output.append(var)
         return tuple(output)
     else:
         dims, = output_dims
-        return Variable(dims, result_data)
+        var = Variable(dims, result_data)
+        if keep_attrs and isinstance(args[0], Variable):
+            var.attrs.update(args[0].attrs)
+        return var
 
 
 def _apply_with_dask_atop(func, args, input_dims, output_dims, signature,
@@ -902,6 +903,7 @@ def apply_ufunc(func, *args, **kwargs):
     variables_ufunc = functools.partial(apply_variable_ufunc, func,
                                         signature=signature,
                                         exclude_dims=exclude_dims,
+                                        keep_attrs=keep_attrs,
                                         dask=dask,
                                         output_dtypes=output_dtypes,
                                         output_sizes=output_sizes)
@@ -930,8 +932,7 @@ def apply_ufunc(func, *args, **kwargs):
         return apply_dataarray_ufunc(variables_ufunc, *args,
                                      signature=signature,
                                      join=join,
-                                     exclude_dims=exclude_dims,
-                                     keep_attrs=keep_attrs)
+                                     exclude_dims=exclude_dims)
     elif any(isinstance(a, Variable) for a in args):
         return variables_ufunc(*args)
     else:


Test patch:
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index db10ee3e..c84ed17b 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -480,12 +480,19 @@ def test_keep_attrs():
 
     a = xr.DataArray([0, 1], [('x', [0, 1])])
     a.attrs['attr'] = 'da'
+    a['x'].attrs['attr'] = 'da_coord'
     b = xr.DataArray([1, 2], [('x', [0, 1])])
 
     actual = add(a, b, keep_attrs=False)
     assert not actual.attrs
     actual = add(a, b, keep_attrs=True)
     assert_identical(actual.attrs, a.attrs)
+    assert_identical(actual['x'].attrs, a['x'].attrs)
+
+    actual = add(a.variable, b.variable, keep_attrs=False)
+    assert not actual.attrs
+    actual = add(a.variable, b.variable, keep_attrs=True)
+    assert_identical(actual.attrs, a.attrs)
 
     a = xr.Dataset({'x': [0, 1]})
     a.attrs['attr'] = 'ds'