Support for dict-like arguments for assign_coords (pydata#3243)

* Issue 3231: Support for dict-like arguments for assign_coords

* Amended docstring to clarify two examples.

* Simplified documentation by not repeating arguments. Formatting.

Author: gwgundersen

doc/whats-new.rst

@@ -76,6 +76,10 @@ Enhancements
 
 Bug fixes
 ~~~~~~~~~
+
+- :py:meth:`~xarray.DataArray.assign_coords` now supports dictionary arguments
+  (:issue:`3231`).
+  By `Gregory Gundersen <https://github.com/gwgundersen>`_.
 - Fix regression introduced in v0.12.2 where ``copy(deep=True)`` would convert
   unicode indices to dtype=object (:issue:`3094`).
   By `Guido Imperiale <https://github.com/crusaderky>`_.

xarray/core/common.py

@@ -341,19 +341,23 @@ def _calc_assign_results(
                 results[k] = v
         return results
 
-    def assign_coords(self, **kwargs):
+    def assign_coords(self, coords=None, **coords_kwargs):
         """Assign new coordinates to this object.
 
         Returns a new object with all the original data in addition to the new
         coordinates.
 
         Parameters
         ----------
-        kwargs : keyword, value pairs
-            keywords are the variables names. If the values are callable, they
-            are computed on this object and assigned to new coordinate
-            variables. If the values are not callable, (e.g. a DataArray,
-            scalar, or array), they are simply assigned.
+        coords : dict, optional
+            A dict with keys which are variables names. If the values are
+            callable, they are computed on this object and assigned to new
+            coordinate variables. If the values are not callable,
+            (e.g. a ``DataArray``, scalar, or array), they are simply assigned.
+
+        **coords_kwargs : keyword, value pairs, optional
+            The keyword arguments form of ``coords``.
+            One of ``coords`` or ``coords_kwargs`` must be provided.
 
         Returns
         -------
@@ -363,7 +367,6 @@ def assign_coords(self, **kwargs):
 
         Examples
         --------
-
         Convert longitude coordinates from 0-359 to -180-179:
 
         >>> da = xr.DataArray(np.random.rand(4),
@@ -380,10 +383,18 @@ def assign_coords(self, **kwargs):
         Coordinates:
           * lon      (lon) int64 -2 -1 0 1
 
+        The function also accepts dictionary arguments:
+
+        >>> da.assign_coords({'lon': (((da.lon + 180) % 360) - 180)})
+        <xarray.DataArray (lon: 4)>
+        array([0.28298 , 0.667347, 0.657938, 0.177683])
+        Coordinates:
+          * lon      (lon) int64 -2 -1 0 1
+
         Notes
         -----
-        Since ``kwargs`` is a dictionary, the order of your arguments may not
-        be preserved, and so the order of the new variables is not well
+        Since ``coords_kwargs`` is a dictionary, the order of your arguments may
+        not be preserved, and so the order of the new variables is not well
         defined. Assigning multiple variables within the same ``assign_coords``
         is possible, but you cannot reference other variables created within
         the same ``assign_coords`` call.
@@ -393,8 +404,9 @@ def assign_coords(self, **kwargs):
         Dataset.assign
         Dataset.swap_dims
         """
+        coords_kwargs = either_dict_or_kwargs(coords, coords_kwargs, "assign_coords")
         data = self.copy(deep=False)
-        results = self._calc_assign_results(kwargs)
+        results = self._calc_assign_results(coords_kwargs)
         data.coords.update(results)
         return data
 

xarray/core/groupby.py

@@ -11,7 +11,13 @@
 from .common import ALL_DIMS, ImplementsArrayReduce, ImplementsDatasetReduce
 from .options import _get_keep_attrs
 from .pycompat import integer_types
-from .utils import hashable, maybe_wrap_array, peek_at, safe_cast_to_index
+from .utils import (
+    hashable,
+    maybe_wrap_array,
+    peek_at,
+    safe_cast_to_index,
+    either_dict_or_kwargs,
+)
 from .variable import IndexVariable, Variable, as_variable
 
 
@@ -507,15 +513,16 @@ def last(self, skipna=None, keep_attrs=None):
         """
         return self._first_or_last(duck_array_ops.last, skipna, keep_attrs)
 
-    def assign_coords(self, **kwargs):
+    def assign_coords(self, coords=None, **coords_kwargs):
         """Assign coordinates by group.
 
         See also
         --------
         Dataset.assign_coords
         Dataset.swap_dims
         """
-        return self.apply(lambda ds: ds.assign_coords(**kwargs))
+        coords_kwargs = either_dict_or_kwargs(coords, coords_kwargs, "assign_coords")
+        return self.apply(lambda ds: ds.assign_coords(**coords_kwargs))
 
 
 def _maybe_reorder(xarray_obj, dim, positions):

xarray/tests/test_dataset.py

@@ -3198,6 +3198,19 @@ def test_assign(self):
         expected = expected.set_coords("c")
         assert_identical(actual, expected)
 
+    def test_assign_coords(self):
+        ds = Dataset()
+
+        actual = ds.assign(x=[0, 1, 2], y=2)
+        actual = actual.assign_coords(x=list("abc"))
+        expected = Dataset({"x": list("abc"), "y": 2})
+        assert_identical(actual, expected)
+
+        actual = ds.assign(x=[0, 1, 2], y=[2, 3])
+        actual = actual.assign_coords({"y": [2.0, 3.0]})
+        expected = ds.assign(x=[0, 1, 2], y=[2.0, 3.0])
+        assert_identical(actual, expected)
+
     def test_assign_attrs(self):
         expected = Dataset(attrs=dict(a=1, b=2))
         new = Dataset()

xarray/tests/test_groupby.py

@@ -189,4 +189,17 @@ def test_da_groupby_quantile():
     assert_identical(expected, actual)
 
 
+def test_da_groupby_assign_coords():
+    actual = xr.DataArray(
+        [[3, 4, 5], [6, 7, 8]], dims=["y", "x"], coords={"y": range(2), "x": range(3)}
+    )
+    actual1 = actual.groupby("x").assign_coords({"y": [-1, -2]})
+    actual2 = actual.groupby("x").assign_coords(y=[-1, -2])
+    expected = xr.DataArray(
+        [[3, 4, 5], [6, 7, 8]], dims=["y", "x"], coords={"y": [-1, -2], "x": range(3)}
+    )
+    assert_identical(expected, actual1)
+    assert_identical(expected, actual2)
+
+
 # TODO: move other groupby tests from test_dataset and test_dataarray over here