add assign_attrs method (#1286)

* add assign_attrs method

* Add assign_attrs doc links.

* Fix docstring formatting.

Author: hsharrison

doc/api.rst

@@ -75,6 +75,7 @@ Dataset contents
    Dataset.copy
    Dataset.assign
    Dataset.assign_coords
+   Dataset.assign_attrs
    Dataset.pipe
    Dataset.merge
    Dataset.rename
@@ -217,6 +218,7 @@ DataArray contents
    :toctree: generated/
 
    DataArray.assign_coords
+   DataArray.assign_attrs
    DataArray.pipe
    DataArray.rename
    DataArray.swap_dims

doc/whats-new.rst

@@ -30,6 +30,11 @@ Enhancements
   specified with `x=` or `y=`, the other dimension is now guessed. By
   `Vincent Noel <https://github.com/vnoel>`_.
 
+Added new method :py:meth:`~Dataset.assign_attrs` to ``DataArray`` and
+``Dataset``, a chained-method compatible implementation of the
+``dict.update`` method on attrs (:issue:`1281`).
+By `Henry S. Harrison <https://hsharrison.github.io>`_.
+
 Bug fixes
 ~~~~~~~~~
 - ``rolling`` now keeps its original dimension order (:issue:`1125`).

xarray/core/common.py

@@ -312,8 +312,10 @@ def _calc_assign_results(self, kwargs):
         return results
 
     def assign_coords(self, **kwargs):
-        """Assign new coordinates to this object, returning a new object
-        with all the original data in addition to the new coordinates.
+        """Assign new coordinates to this object.
+
+        Returns a new object with all the original data in addition to the new
+        coordinates.
 
         Parameters
         ----------
@@ -346,6 +348,29 @@ def assign_coords(self, **kwargs):
         data.coords.update(results)
         return data
 
+    def assign_attrs(self, *args, **kwargs):
+        """Assign new attrs to this object.
+
+        Returns a new object equivalent to self.attrs.update(*args, **kwargs).
+
+        Parameters
+        ----------
+        args : positional arguments passed into ``attrs.update``.
+        kwargs : keyword arguments passed into ``attrs.update``.
+
+        Returns
+        -------
+        assigned : same type as caller
+            A new object with the new attrs in addition to the existing data.
+
+        See also
+        --------
+        Dataset.assign
+        """
+        out = self.copy(deep=False)
+        out.attrs.update(*args, **kwargs)
+        return out
+
     def pipe(self, func, *args, **kwargs):
         """
         Apply func(self, *args, **kwargs)

xarray/tests/test_dataarray.py

@@ -1352,6 +1352,20 @@ def test_reduce_keep_attrs(self):
         self.assertEqual(len(vm.attrs), len(self.attrs))
         self.assertEqual(vm.attrs, self.attrs)
 
+    def test_assign_attrs(self):
+        expected = DataArray([], attrs=dict(a=1, b=2))
+        expected.attrs['a'] = 1
+        expected.attrs['b'] = 2
+        new = DataArray([])
+        actual = DataArray([]).assign_attrs(a=1, b=2)
+        self.assertDatasetIdentical(actual, expected)
+        self.assertEqual(new.attrs, {})
+
+        expected.attrs['c'] = 3
+        new_actual = actual.assign_attrs({'c': 3})
+        self.assertDatasetIdentical(new_actual, expected)
+        self.assertEqual(actual.attrs, {'a': 1, 'b': 2})
+
     def test_fillna(self):
         a = DataArray([np.nan, 1, np.nan, 3], coords={'x': range(4)}, dims='x')
         actual = a.fillna(-1)

xarray/tests/test_dataset.py

@@ -1848,6 +1848,18 @@ def test_assign(self):
         expected = expected.set_coords('c')
         self.assertDatasetIdentical(actual, expected)
 
+    def test_assign_attrs(self):
+        expected = Dataset(attrs=dict(a=1, b=2))
+        new = Dataset()
+        actual = new.assign_attrs(a=1, b=2)
+        self.assertDatasetIdentical(actual, expected)
+        self.assertEqual(new.attrs, {})
+
+        expected.attrs['c'] = 3
+        new_actual = actual.assign_attrs({'c': 3})
+        self.assertDatasetIdentical(new_actual, expected)
+        self.assertEqual(actual.attrs, dict(a=1, b=2))
+
     def test_assign_multiindex_level(self):
         data = create_test_multiindex()
         with self.assertRaisesRegexp(ValueError, 'conflicting MultiIndex'):