Add Coordinates.assign() method

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

@@ -23,6 +23,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,7 @@
 )
 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
+from xarray.core.utils import Frozen, ReprObject, either_dict_or_kwargs
 from xarray.core.variable import Variable, as_variable, calculate_dimensions
 
 if TYPE_CHECKING:
@@ -472,6 +472,35 @@ 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.
+
+        """
+        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

@@ -112,6 +112,21 @@ def test_identical(self):
         assert coords.identical(coords)
         assert not coords.identical("no_a_coords")
 
+    def test_assign(self) -> None:
+        _ds = Dataset(coords={"x": [0, 1, 2]})
+        coords = Coordinates(coords=_ds.coords, indexes=_ds.xindexes)
+
+        _ds_expected = Dataset(coords={"x": [0, 1, 2], "y": [3, 4]})
+        expected = Coordinates(
+            coords=_ds_expected.coords, indexes=_ds_expected.xindexes
+        )
+
+        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()