diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 13fe0a791bb..e2008687b24 100644 --- a/xarray/core/coordinates.py +++ b/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,31 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset: variables=coords, coord_names=coord_names, indexes=indexes ) + def __or__(self, other: Self) -> Self: + """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. + """ + return self.merge(other).coords + def __setitem__(self, key: Hashable, value: Any) -> None: self.update({key: value}) @@ -719,6 +744,78 @@ def copy( ), ) + def drop_vars( + self, + names: str | Iterable[Hashable] | Callable[[Self], 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 self.to_dataset().drop_vars(names, 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 Dataset. + **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 self.to_dataset().rename_dims(dims_dict, **dims).coords + + def rename_vars( + self, + name_dict: Mapping[Any, Hashable] | None = None, + **names: Hashable, + ) -> Self: + """Returns a new object with renamed variables including coordinates + + 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 : Dataset + Dataset with renamed variables including coordinates + """ + return self.to_dataset().rename_vars(name_dict, **names) + class DatasetCoordinates(Coordinates): """Dictionary like container for Dataset coordinates (variables + indexes). diff --git a/xarray/tests/test_coordinates.py b/xarray/tests/test_coordinates.py index b4ba922203a..3f95acaf005 100644 --- a/xarray/tests/test_coordinates.py +++ b/xarray/tests/test_coordinates.py @@ -208,3 +208,63 @@ 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 set(actual.variables) == {"a", "y"} + + actual = coords.drop_vars(["x", "y"]) + assert set(actual.variables) == {"a"} + + def test_rename_dims(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.rename_dims({"x": "X"}) + assert set(actual.dims) == {"X", "y"} + assert set(actual.variables) == {"a", "x", "y"} + + actual = coords.rename_dims({"x": "u", "y": "v"}) + assert set(actual.dims) == {"u", "v"} + assert set(actual.variables) == {"a", "x", "y"} + + def test_rename_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.rename_vars({"x": "X"}) + assert set(actual.dims) == {"x", "y"} + assert set(actual.variables) == {"a", "X", "y"} + + actual = coords.rename_vars({"x": "u", "y": "v"}) + assert set(actual.dims) == {"x", "y"} + assert set(actual.variables) == {"a", "u", "v"} + + def test_operator_merge(self): + 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)