Add keep_attrs to reduction methods

test/test_data_array.py

@@ -2,6 +2,7 @@
 import pandas as pd
 from copy import deepcopy
 from textwrap import dedent
+from collections import OrderedDict
 
 from xray import Dataset, DataArray, Variable, align
 from xray.pycompat import iteritems
@@ -10,8 +11,10 @@
 
 class TestDataArray(TestCase):
     def setUp(self):
+        self._attrs = {'attr1': 'value1', 'attr2': 2929}
         self.x = np.random.random((10, 20))
         self.v = Variable(['x', 'y'], self.x)
+        self.va = Variable(['x', 'y'], self.x, self._attrs)
         self.ds = Dataset({'foo': self.v})
         self.dv = self.ds['foo']
 
@@ -262,6 +265,17 @@ def test_reduce(self):
         # needs more...
         # should check which extra dimensions are dropped
 
+    def test_reduce_keep_attrs(self):
+        # Test dropped attrs
+        vm = self.va.mean()
+        self.assertEqual(len(vm.attrs), 0)
+        self.assertEqual(vm.attrs, OrderedDict())
+
+        # Test kept attrs
+        vm = self.va.mean(keep_attrs=True)
+        self.assertEqual(len(vm.attrs), len(self._attrs))
+        self.assertEqual(vm.attrs, self._attrs)
+
     def test_unselect(self):
         with self.assertRaisesRegexp(ValueError, 'cannot unselect the name'):
             self.dv.unselect('foo')

test/test_dataset.py

@@ -698,3 +698,20 @@ def test_reduce_non_numeric(self):
         self.assertDatasetEqual(data1.mean(), data2.mean())
         self.assertDatasetEqual(data1.mean(dimension='dim1'),
                                 data2.mean(dimension='dim1'))
+
+    def test_reduce_keep_attrs(self):
+        data = create_test_data()
+        _attrs = {'attr1': 'value1', 'attr2': 2929}
+
+        attrs = OrderedDict(_attrs)
+        data.attrs = attrs
+
+        # Test dropped attrs
+        ds = data.mean()
+        self.assertEqual(len(ds.attrs), 0)
+        self.assertEqual(ds.attrs, OrderedDict())
+
+        # Test kept attrs
+        ds = data.mean(keep_attrs=True)
+        self.assertEqual(len(ds.attrs), len(_attrs))
+        self.assertTrue(ds.attrs, attrs)

test/test_variable.py

@@ -530,6 +530,21 @@ def test_reduce(self):
         with self.assertRaisesRegexp(ValueError, 'cannot supply both'):
             v.mean(dimension='x', axis=0)
 
+    def test_reduce_keep_attrs(self):
+        _attrs = {'units': 'test', 'long_name': 'testing'}
+
+        v = Variable(['x', 'y'], self.d, _attrs)
+
+        # Test dropped attrs
+        vm = v.mean()
+        self.assertEqual(len(vm.attrs), 0)
+        self.assertEqual(vm.attrs, OrderedDict())
+
+        # Test kept attrs
+        vm = v.mean(keep_attrs=True)
+        self.assertEqual(len(vm.attrs), len(_attrs))
+        self.assertEqual(vm.attrs, _attrs)
+
 
 class TestCoordinate(TestCase, VariableSubclassTestCases):
     cls = staticmethod(Coordinate)

xray/common.py

@@ -5,8 +5,8 @@
 class ImplementsReduce(object):
     @classmethod
     def _reduce_method(cls, f, name=None, module=None):
-        def func(self, dimension=None, axis=None, **kwargs):
-            return self.reduce(f, dimension, axis, **kwargs)
+        def func(self, dimension=None, axis=None, keep_attrs=False, **kwargs):
+            return self.reduce(f, dimension, axis, keep_attrs, **kwargs)
         if name is None:
             name = f.__name__
         func.__name__ = name
@@ -96,6 +96,10 @@ def _get_axis_num(self, dim):
             and 'axis' arguments can be supplied. If neither are supplied, then
             `{name}` is calculated over the flattened array (by calling
             `{name}(x)` without an axis argument).
+        keep_attrs : bool, optional
+            If True, the variable's attributes (`attrs`) will be copied from
+            the original object to the new one.  If False (default), the new
+            object will be returned without attributes.
         **kwargs : dict
             Additional keyword arguments passed on to `{name}`.
 

xray/data_array.py

@@ -465,7 +465,8 @@ def squeeze(self, dimension=None):
         ds = self.dataset.squeeze(dimension)
         return ds[self.name]
 
-    def reduce(self, func, dimension=None, axis=None, **kwargs):
+    def reduce(self, func, dimension=None, axis=None, keep_attrs=False,
+               **kwargs):
         """Reduce this array by applying `func` along some dimension(s).
 
         Parameters
@@ -481,6 +482,10 @@ def reduce(self, func, dimension=None, axis=None, **kwargs):
             'dimension' and 'axis' arguments can be supplied. If neither are
             supplied, then the reduction is calculated over the flattened array
             (by calling `f(x)` without an axis argument).
+        keep_attrs : bool, optional
+            If True, the variable's attributes (`attrs`) will be copied from
+            the original object to the new one.  If False (default), the new
+            object will be returned without attributes.
         **kwargs : dict
             Additional keyword arguments passed on to `func`.
 
@@ -490,7 +495,7 @@ def reduce(self, func, dimension=None, axis=None, **kwargs):
             DataArray with this object's array replaced with an array with
             summarized data and the indicated dimension(s) removed.
         """
-        var = self.variable.reduce(func, dimension, axis, **kwargs)
+        var = self.variable.reduce(func, dimension, axis, keep_attrs, **kwargs)
         drop = set(self.dimensions) - set(var.dimensions)
         # For now, take an aggressive strategy of removing all variables
         # associated with any dropped dimensions
@@ -499,6 +504,10 @@ def reduce(self, func, dimension=None, axis=None, **kwargs):
                  if any(dim in drop for dim in v.dimensions)}
         ds = self.dataset.unselect(*drop)
         ds[self.name] = var
+
+        if keep_attrs:
+            ds.attrs = self.dataset.attrs
+
         return ds[self.name]
 
     @classmethod

xray/dataset.py

@@ -983,6 +983,10 @@ def squeeze(self, dimension=None):
         dimension : str or sequence of str, optional
             Dimension(s) over which to apply `func`.  By default `func` is
             applied over all dimensions.
+        keep_attrs : bool, optional
+            If True, the datasets's attributes (`attrs`) will be copied from
+            the original object to the new one.  If False (default), the new
+            object will be returned without attributes.
         **kwargs : dict
             Additional keyword arguments passed on to `{name}`.
 
@@ -995,8 +999,8 @@ def squeeze(self, dimension=None):
 
     @classmethod
     def _reduce_method(cls, f, name=None, module=None):
-        def func(self, dimension=None, **kwargs):
-            return self.reduce(f, dimension, **kwargs)
+        def func(self, dimension=None, keep_attrs=False, **kwargs):
+            return self.reduce(f, dimension, keep_attrs, **kwargs)
         if name is None:
             name = f.__name__
         func.__name__ = name
@@ -1005,7 +1009,7 @@ def func(self, dimension=None, **kwargs):
             cls=cls.__name__)
         return func
 
-    def reduce(self, func, dimension=None, **kwargs):
+    def reduce(self, func, dimension=None, keep_attrs=False, **kwargs):
         """Reduce this dataset by applying `func` along some dimension(s).
 
         Parameters
@@ -1017,6 +1021,10 @@ def reduce(self, func, dimension=None, **kwargs):
         dimension : str or sequence of str, optional
             Dimension(s) over which to apply `func`.  By default `func` is
             applied over all dimensions.
+        keep_attrs : bool, optional
+            If True, the datasets's attributes (`attrs`) will be copied from
+            the original object to the new one.  If False (default), the new
+            object will be returned without attributes.
         **kwargs : dict
             Additional keyword arguments passed on to `func`.
 
@@ -1052,7 +1060,10 @@ def reduce(self, func, dimension=None, **kwargs):
                         pass
             else:
                 variables[name] = var
-        return Dataset(variables=variables)
+
+        attrs = self.attrs if keep_attrs else {}
+
+        return Dataset(variables=variables, attributes=attrs)
 
     @classmethod
     def concat(cls, datasets, dimension='concat_dimension', indexers=None,

xray/variable.py

@@ -457,7 +457,8 @@ def squeeze(self, dimension=None):
         dimensions = dict(zip(self.dimensions, self.shape))
         return utils.squeeze(self, dimensions, dimension)
 
-    def reduce(self, func, dimension=None, axis=None, **kwargs):
+    def reduce(self, func, dimension=None, axis=None, keep_attrs=False,
+               **kwargs):
         """Reduce this array by applying `func` along some dimension(s).
 
         Parameters
@@ -473,6 +474,10 @@ def reduce(self, func, dimension=None, axis=None, **kwargs):
             and 'axis' arguments can be supplied. If neither are supplied, then
             the reduction is calculated over the flattened array (by calling
             `func(x)` without an axis argument).
+        keep_attrs : bool, optional
+            If True, the variable's attributes (`attrs`) will be copied from
+            the original object to the new one.  If False (default), the new
+            object will be returned without attributes.
         **kwargs : dict
             Additional keyword arguments passed on to `func`.
 
@@ -482,6 +487,7 @@ def reduce(self, func, dimension=None, axis=None, **kwargs):
             Array with summarized data and the indicated dimension(s)
             removed.
         """
+
         if dimension is not None and axis is not None:
             raise ValueError("cannot supply both 'axis' and 'dimension' "
                              "arguments")
@@ -495,7 +501,9 @@ def reduce(self, func, dimension=None, axis=None, **kwargs):
         dims = [dim for n, dim in enumerate(self.dimensions)
                 if n not in removed_axes]
 
-        return Variable(dims, data)
+        attrs = self.attrs if keep_attrs else {}
+
+        return Variable(dims, data, attributes=attrs)
 
     @classmethod
     def concat(cls, variables, dimension='stacked_dimension',