diff --git a/test/test_data_array.py b/test/test_data_array.py index cbbab21774b..edac6926039 100644 --- a/test/test_data_array.py +++ b/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') diff --git a/test/test_dataset.py b/test/test_dataset.py index 9f259d81d45..8aa602bd76f 100644 --- a/test/test_dataset.py +++ b/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) diff --git a/test/test_variable.py b/test/test_variable.py index 632db73a0a0..30250378d21 100644 --- a/test/test_variable.py +++ b/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) diff --git a/xray/common.py b/xray/common.py index 3575d9d4115..549dee1102b 100644 --- a/xray/common.py +++ b/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}`. diff --git a/xray/data_array.py b/xray/data_array.py index 2a393cee75e..dff43651689 100644 --- a/xray/data_array.py +++ b/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 diff --git a/xray/dataset.py b/xray/dataset.py index 27a80c4cd0f..76e3f6761c1 100644 --- a/xray/dataset.py +++ b/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, diff --git a/xray/variable.py b/xray/variable.py index 97b43e0ebb5..402997731f3 100644 --- a/xray/variable.py +++ b/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',