Add Coordinates.assign() method (#8102)

* add Coordinates.assign method

* update what's new

* test typing

Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>

* add docstring examples

* fix doctests

* fix doctests (spaces)

* doctests typo again

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix assign test after merging main

---------

Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

doc/api-hidden.rst

@@ -22,6 +22,7 @@
    Coordinates.to_dataset
    Coordinates.to_index
    Coordinates.update
+   Coordinates.assign
    Coordinates.merge
    Coordinates.copy
    Coordinates.equals
@@ -39,8 +40,9 @@
    core.coordinates.DatasetCoordinates.to_dataset
    core.coordinates.DatasetCoordinates.to_index
    core.coordinates.DatasetCoordinates.update
+   core.coordinates.DatasetCoordinates.assign
    core.coordinates.DatasetCoordinates.merge
-   core.coordinates.DataArrayCoordinates.copy
+   core.coordinates.DatasetCoordinates.copy
    core.coordinates.DatasetCoordinates.equals
    core.coordinates.DatasetCoordinates.identical
 
@@ -79,6 +81,7 @@
    core.coordinates.DataArrayCoordinates.to_dataset
    core.coordinates.DataArrayCoordinates.to_index
    core.coordinates.DataArrayCoordinates.update
+   core.coordinates.DataArrayCoordinates.assign
    core.coordinates.DataArrayCoordinates.merge
    core.coordinates.DataArrayCoordinates.copy
    core.coordinates.DataArrayCoordinates.equals

doc/whats-new.rst

@@ -22,6 +22,10 @@ v2023.08.1 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Added the :py:meth:`Coordinates.assign` method that can be used to combine
+  different collections of coordinates prior to assign them to a Dataset or
+  DataArray (:pull:`8102`) at once.
+  By `Benoît Bovy <https://github.com/benbovy>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/core/coordinates.py

@@ -24,7 +24,12 @@
 )
 from xarray.core.merge import merge_coordinates_without_align, merge_coords
 from xarray.core.types import Self, T_DataArray
-from xarray.core.utils import Frozen, ReprObject, emit_user_level_warning
+from xarray.core.utils import (
+    Frozen,
+    ReprObject,
+    either_dict_or_kwargs,
+    emit_user_level_warning,
+)
 from xarray.core.variable import Variable, as_variable, calculate_dimensions
 
 if TYPE_CHECKING:
@@ -560,6 +565,53 @@ def update(self, other: Mapping[Any, Any]) -> None:
 
         self._update_coords(coords, indexes)
 
+    def assign(
+        self, coords: Mapping | None = None, **coords_kwargs: Any
+    ) -> Coordinates:
+        """Assign new coordinates (and indexes) to a Coordinates object, returning
+        a new object with all the original coordinates in addition to the new ones.
+
+        Parameters
+        ----------
+        coords : :class:`Coordinates` or mapping of hashable to Any
+            Mapping from coordinate names to the new values. If a ``Coordinates``
+            object is passed, its indexes are assigned in the returned object.
+            Otherwise, a default (pandas) index is created for each dimension
+            coordinate found in the mapping.
+        **coords_kwargs
+            The keyword arguments form of ``coords``.
+            One of ``coords`` or ``coords_kwargs`` must be provided.
+
+        Returns
+        -------
+        new_coords : Coordinates
+            A new Coordinates object with the new coordinates (and indexes)
+            in addition to all the existing coordinates.
+
+        Examples
+        --------
+        >>> coords = xr.Coordinates()
+        >>> coords
+        Coordinates:
+            *empty*
+
+        >>> coords.assign(x=[1, 2])
+        Coordinates:
+          * x        (x) int64 1 2
+
+        >>> midx = pd.MultiIndex.from_product([["a", "b"], [0, 1]])
+        >>> coords.assign(xr.Coordinates.from_pandas_multiindex(midx, "y"))
+        Coordinates:
+          * y          (y) object MultiIndex
+          * y_level_0  (y) object 'a' 'a' 'b' 'b'
+          * y_level_1  (y) int64 0 1 0 1
+
+        """
+        coords = either_dict_or_kwargs(coords, coords_kwargs, "assign")
+        new_coords = self.copy()
+        new_coords.update(coords)
+        return new_coords
+
     def _overwrite_indexes(
         self,
         indexes: Mapping[Any, Index],

xarray/tests/test_coordinates.py

@@ -124,6 +124,16 @@ def test_identical(self):
         assert coords.identical(coords)
         assert not coords.identical("not_a_coords")
 
+    def test_assign(self) -> None:
+        coords = Coordinates(coords={"x": [0, 1, 2]})
+        expected = Coordinates(coords={"x": [0, 1, 2], "y": [3, 4]})
+
+        actual = coords.assign(y=[3, 4])
+        assert_identical(actual, expected)
+
+        actual = coords.assign({"y": [3, 4]})
+        assert_identical(actual, expected)
+
     def test_copy(self) -> None:
         no_index_coords = Coordinates({"foo": ("x", [1, 2, 3])})
         copied = no_index_coords.copy()