modification methods on Coordinates (pydata#10318)

* implement  `Coordinates` methods modifying coords

* allow using the binary-or operator (`|`) for merging

* tests for `drop_vars`

* tests for `rename_dims` and `rename_vars`

* tests for the merge operator

* Apply suggestions from code review

Co-authored-by: Benoit Bovy <[email protected]>

* attempt to fix the typing

* make sure we always return a `Coordinates` object

* replace docstring by a reference to `Coordinates.merge`

* changelog

* create docs pages

* add back the docstring for `__or__`

* add `drop_dims`

* typing

* undo a bad merge [skip-ci]

* document `drop_dims` [skip-ci]

---------

Co-authored-by: Benoit Bovy <[email protected]>
Co-authored-by: Deepak Cherian <[email protected]>

Author: 3 people

doc/api/coordinates.rst

@@ -38,6 +38,7 @@ and values given by ``DataArray`` objects.
    Coordinates.__getitem__
    Coordinates.__setitem__
    Coordinates.__delitem__
+   Coordinates.__or__
    Coordinates.update
    Coordinates.get
    Coordinates.items
@@ -53,8 +54,12 @@ Coordinates contents
    Coordinates.to_dataset
    Coordinates.to_index
    Coordinates.assign
+   Coordinates.drop_dims
+   Coordinates.drop_vars
    Coordinates.merge
    Coordinates.copy
+   Coordinates.rename_vars
+   Coordinates.rename_dims
 
 Comparisons
 -----------

doc/whats-new.rst

@@ -12,7 +12,8 @@ v2025.08.1 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
-
+- Add convenience methods to :py:class:`~xarray.Coordinates` (:pull:`10318`)
+  By `Justus Magin <https://github.com/keewis>`_.
 - Added :py:func:`load_datatree` for loading ``DataTree`` objects into memory
   from disk. It has the same relationship to :py:func:`open_datatree`, as
   :py:func:`load_dataset` has to :py:func:`open_dataset`.

xarray/core/coordinates.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Hashable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from typing import (
     TYPE_CHECKING,
@@ -21,7 +21,7 @@
     assert_no_index_corrupted,
     create_default_index_implicit,
 )
-from xarray.core.types import DataVars, Self, T_DataArray, T_Xarray
+from xarray.core.types import DataVars, ErrorOptions, Self, T_DataArray, T_Xarray
 from xarray.core.utils import (
     Frozen,
     ReprObject,
@@ -561,6 +561,35 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset:
             variables=coords, coord_names=coord_names, indexes=indexes
         )
 
+    def __or__(self, other: Mapping[Any, Any] | None) -> Coordinates:
+        """Merge two sets of coordinates to create a new Coordinates object
+
+        The method implements the logic used for joining coordinates in the
+        result of a binary operation performed on xarray objects:
+
+        - If two index coordinates conflict (are not equal), an exception is
+          raised. You must align your data before passing it to this method.
+        - If an index coordinate and a non-index coordinate conflict, the non-
+          index coordinate is dropped.
+        - If two non-index coordinates conflict, both are dropped.
+
+        Parameters
+        ----------
+        other : dict-like, optional
+            A :py:class:`Coordinates` object or any mapping that can be turned
+            into coordinates.
+
+        Returns
+        -------
+        merged : Coordinates
+            A new Coordinates object with merged coordinates.
+
+        See Also
+        --------
+        Coordinates.merge
+        """
+        return self.merge(other).coords
+
     def __setitem__(self, key: Hashable, value: Any) -> None:
         self.update({key: value})
 
@@ -719,6 +748,108 @@ def copy(
             ),
         )
 
+    def drop_vars(
+        self,
+        names: str
+        | Iterable[Hashable]
+        | Callable[
+            [Coordinates | Dataset | DataArray | DataTree],
+            str | Iterable[Hashable],
+        ],
+        *,
+        errors: ErrorOptions = "raise",
+    ) -> Self:
+        """Drop variables from this Coordinates object.
+
+        Note that indexes that depend on these variables will also be dropped.
+
+        Parameters
+        ----------
+        names : hashable or iterable or callable
+            Name(s) of variables to drop. If a callable, this is object is passed as its
+            only argument and its result is used.
+        errors : {"raise", "ignore"}, default: "raise"
+            Error treatment.
+
+            - ``'raise'``: raises a :py:class:`ValueError` error if any of the variable
+              passed are not in the dataset
+            - ``'ignore'``: any given names that are in the dataset are dropped and no
+              error is raised.
+        """
+        return cast(Self, self.to_dataset().drop_vars(names, errors=errors).coords)
+
+    def drop_dims(
+        self,
+        drop_dims: str | Iterable[Hashable],
+        *,
+        errors: ErrorOptions = "raise",
+    ) -> Self:
+        """Drop dimensions and associated variables from this dataset.
+
+        Parameters
+        ----------
+        drop_dims : str or Iterable of Hashable
+            Dimension or dimensions to drop.
+        errors : {"raise", "ignore"}, default: "raise"
+            If 'raise', raises a ValueError error if any of the
+            dimensions passed are not in the dataset. If 'ignore', any given
+            dimensions that are in the dataset are dropped and no error is raised.
+
+        Returns
+        -------
+        obj : Coordinates
+            Coordinates object without the given dimensions (or any coordinates
+            containing those dimensions).
+        """
+        return cast(Self, self.to_dataset().drop_dims(drop_dims, errors=errors).coords)
+
+    def rename_dims(
+        self,
+        dims_dict: Mapping[Any, Hashable] | None = None,
+        **dims: Hashable,
+    ) -> Self:
+        """Returns a new object with renamed dimensions only.
+
+        Parameters
+        ----------
+        dims_dict : dict-like, optional
+            Dictionary whose keys are current dimension names and
+            whose values are the desired names. The desired names must
+            not be the name of an existing dimension or Variable in the Coordinates.
+        **dims : optional
+            Keyword form of ``dims_dict``.
+            One of dims_dict or dims must be provided.
+
+        Returns
+        -------
+        renamed : Coordinates
+            Coordinates object with renamed dimensions.
+        """
+        return cast(Self, self.to_dataset().rename_dims(dims_dict, **dims).coords)
+
+    def rename_vars(
+        self,
+        name_dict: Mapping[Any, Hashable] | None = None,
+        **names: Hashable,
+    ) -> Coordinates:
+        """Returns a new object with renamed variables.
+
+        Parameters
+        ----------
+        name_dict : dict-like, optional
+            Dictionary whose keys are current variable or coordinate names and
+            whose values are the desired names.
+        **names : optional
+            Keyword form of ``name_dict``.
+            One of name_dict or names must be provided.
+
+        Returns
+        -------
+        renamed : Coordinates
+            Coordinates object with renamed variables
+        """
+        return cast(Self, self.to_dataset().rename_vars(name_dict, **names).coords)
+
 
 class DatasetCoordinates(Coordinates):
     """Dictionary like container for Dataset coordinates (variables + indexes).

xarray/tests/test_coordinates.py

@@ -208,3 +208,87 @@ def test_dataset_from_coords_with_multidim_var_same_name(self):
         coords = Coordinates(coords={"x": var}, indexes={})
         ds = Dataset(coords=coords)
         assert ds.coords["x"].dims == ("x", "y")
+
+    def test_drop_vars(self):
+        coords = Coordinates(
+            coords={
+                "x": Variable("x", range(3)),
+                "y": Variable("y", list("ab")),
+                "a": Variable(["x", "y"], np.arange(6).reshape(3, 2)),
+            },
+            indexes={},
+        )
+
+        actual = coords.drop_vars("x")
+        assert isinstance(actual, Coordinates)
+        assert set(actual.variables) == {"a", "y"}
+
+        actual = coords.drop_vars(["x", "y"])
+        assert isinstance(actual, Coordinates)
+        assert set(actual.variables) == {"a"}
+
+    def test_drop_dims(self) -> None:
+        coords = Coordinates(
+            coords={
+                "x": Variable("x", range(3)),
+                "y": Variable("y", list("ab")),
+                "a": Variable(["x", "y"], np.arange(6).reshape(3, 2)),
+            },
+            indexes={},
+        )
+
+        actual = coords.drop_dims("x")
+        assert isinstance(actual, Coordinates)
+        assert set(actual.variables) == {"y"}
+
+        actual = coords.drop_dims(["x", "y"])
+        assert isinstance(actual, Coordinates)
+        assert set(actual.variables) == set()
+
+    def test_rename_dims(self) -> None:
+        coords = Coordinates(
+            coords={
+                "x": Variable("x", range(3)),
+                "y": Variable("y", list("ab")),
+                "a": Variable(["x", "y"], np.arange(6).reshape(3, 2)),
+            },
+            indexes={},
+        )
+
+        actual = coords.rename_dims({"x": "X"})
+        assert isinstance(actual, Coordinates)
+        assert set(actual.dims) == {"X", "y"}
+        assert set(actual.variables) == {"a", "x", "y"}
+
+        actual = coords.rename_dims({"x": "u", "y": "v"})
+        assert isinstance(actual, Coordinates)
+        assert set(actual.dims) == {"u", "v"}
+        assert set(actual.variables) == {"a", "x", "y"}
+
+    def test_rename_vars(self) -> None:
+        coords = Coordinates(
+            coords={
+                "x": Variable("x", range(3)),
+                "y": Variable("y", list("ab")),
+                "a": Variable(["x", "y"], np.arange(6).reshape(3, 2)),
+            },
+            indexes={},
+        )
+
+        actual = coords.rename_vars({"x": "X"})
+        assert isinstance(actual, Coordinates)
+        assert set(actual.dims) == {"x", "y"}
+        assert set(actual.variables) == {"a", "X", "y"}
+
+        actual = coords.rename_vars({"x": "u", "y": "v"})
+        assert isinstance(actual, Coordinates)
+        assert set(actual.dims) == {"x", "y"}
+        assert set(actual.variables) == {"a", "u", "v"}
+
+    def test_operator_merge(self) -> None:
+        coords1 = Coordinates({"x": ("x", [0, 1, 2])})
+        coords2 = Coordinates({"y": ("y", [3, 4, 5])})
+        expected = Dataset(coords={"x": [0, 1, 2], "y": [3, 4, 5]})
+
+        actual = coords1 | coords2
+        assert_identical(Dataset(coords=actual), expected)