From eb6dc62b6f398e0af5be295982d56b26fb57d059 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 13 May 2025 17:22:42 +0200 Subject: [PATCH 01/12] implement `Coordinates` methods modifying coords --- xarray/core/coordinates.py | 76 +++++++++++++++++++++++++++++++++++++- 1 file changed, 74 insertions(+), 2 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 13fe0a791bb..c6c9eb0d15a 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, @@ -719,6 +719,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). From 8a18a52897483de3d821ef55d2ff7a09a4e8b5ac Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 13 May 2025 18:03:51 +0200 Subject: [PATCH 02/12] allow using the binary-or operator (`|`) for merging --- xarray/core/coordinates.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index c6c9eb0d15a..e2008687b24 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -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}) From f8b7f57548008fea467788fa45e468b399451416 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 13 May 2025 18:09:34 +0200 Subject: [PATCH 03/12] tests for `drop_vars` --- xarray/tests/test_coordinates.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/xarray/tests/test_coordinates.py b/xarray/tests/test_coordinates.py index b4ba922203a..d1e07dca12e 100644 --- a/xarray/tests/test_coordinates.py +++ b/xarray/tests/test_coordinates.py @@ -208,3 +208,19 @@ 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"} From 4626ce5a6e0f7e88f73543598feb347a0838ad04 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 13 May 2025 18:36:58 +0200 Subject: [PATCH 04/12] tests for `rename_dims` and `rename_vars` --- xarray/tests/test_coordinates.py | 36 ++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/xarray/tests/test_coordinates.py b/xarray/tests/test_coordinates.py index d1e07dca12e..def2235029b 100644 --- a/xarray/tests/test_coordinates.py +++ b/xarray/tests/test_coordinates.py @@ -224,3 +224,39 @@ def test_drop_vars(self): 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"} From a370ed91b332fcd196cda2afd27496be040ceec0 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 13 May 2025 18:39:06 +0200 Subject: [PATCH 05/12] tests for the merge operator --- xarray/tests/test_coordinates.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/xarray/tests/test_coordinates.py b/xarray/tests/test_coordinates.py index def2235029b..3f95acaf005 100644 --- a/xarray/tests/test_coordinates.py +++ b/xarray/tests/test_coordinates.py @@ -260,3 +260,11 @@ def test_rename_vars(self): 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) From c909ca9d991eca59471910580689dfd2759a95fe Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 28 Jun 2025 23:57:26 +0200 Subject: [PATCH 06/12] Apply suggestions from code review Co-authored-by: Benoit Bovy --- xarray/core/coordinates.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index e2008687b24..8458a012772 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -781,7 +781,7 @@ def rename_dims( 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. + 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. @@ -798,7 +798,7 @@ def rename_vars( name_dict: Mapping[Any, Hashable] | None = None, **names: Hashable, ) -> Self: - """Returns a new object with renamed variables including coordinates + """Returns a new object with renamed variables. Parameters ---------- @@ -811,8 +811,8 @@ def rename_vars( Returns ------- - renamed : Dataset - Dataset with renamed variables including coordinates + renamed : Coordinates + Coordinates object with renamed variables """ return self.to_dataset().rename_vars(name_dict, **names) From 22c89e35531d27f9500636aa518d77a3355d2918 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 29 Jun 2025 10:56:58 +0200 Subject: [PATCH 07/12] attempt to fix the typing --- xarray/core/coordinates.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 4adcb2324da..fb752be8b4f 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -561,7 +561,7 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset: variables=coords, coord_names=coord_names, indexes=indexes ) - def __or__(self, other: Self) -> Self: + 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 @@ -746,7 +746,12 @@ def copy( def drop_vars( self, - names: str | Iterable[Hashable] | Callable[[Self], str | Iterable[Hashable]], + names: str + | Iterable[Hashable] + | Callable[ + [Coordinates | Dataset | DataArray | DataTree], + str | Iterable[Hashable], + ], *, errors: ErrorOptions = "raise", ) -> Self: @@ -767,13 +772,13 @@ def drop_vars( - ``'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 + return cast(Self, self.to_dataset().drop_vars(names, errors=errors).coords) def rename_dims( self, dims_dict: Mapping[Any, Hashable] | None = None, **dims: Hashable, - ) -> Self: + ) -> Coordinates: """Returns a new object with renamed dimensions only. Parameters @@ -797,7 +802,7 @@ def rename_vars( self, name_dict: Mapping[Any, Hashable] | None = None, **names: Hashable, - ) -> Self: + ) -> Coordinates: """Returns a new object with renamed variables. Parameters From 3b9e3e545f7caa63207dff0611a9982078fb5afe Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 29 Jun 2025 10:57:16 +0200 Subject: [PATCH 08/12] make sure we always return a `Coordinates` object --- xarray/core/coordinates.py | 2 +- xarray/tests/test_coordinates.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index fb752be8b4f..90d4fa93318 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -819,7 +819,7 @@ def rename_vars( renamed : Coordinates Coordinates object with renamed variables """ - return self.to_dataset().rename_vars(name_dict, **names) + return self.to_dataset().rename_vars(name_dict, **names).coords class DatasetCoordinates(Coordinates): diff --git a/xarray/tests/test_coordinates.py b/xarray/tests/test_coordinates.py index 3f95acaf005..f8ad096fd10 100644 --- a/xarray/tests/test_coordinates.py +++ b/xarray/tests/test_coordinates.py @@ -220,9 +220,11 @@ def test_drop_vars(self): ) 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_rename_dims(self): @@ -236,10 +238,12 @@ def test_rename_dims(self): ) 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"} @@ -254,10 +258,12 @@ def test_rename_vars(self): ) 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"} From e53b84866a7ce59c8c442bbf62751a0649ff3955 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 29 Jun 2025 10:59:05 +0200 Subject: [PATCH 09/12] replace docstring by a reference to `Coordinates.merge` --- xarray/core/coordinates.py | 22 +++------------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 90d4fa93318..ccb9293b797 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -564,25 +564,9 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset: 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 From 7dbd26a0e1b9fc5b35318ddc6577a924de434f29 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 29 Jun 2025 11:06:36 +0200 Subject: [PATCH 10/12] changelog --- doc/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 6db780484bd..35ec4e26029 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -14,6 +14,8 @@ New Features ~~~~~~~~~~~~ - Expose :py:class:`~xarray.indexes.RangeIndex`, and :py:class:`~xarray.indexes.CoordinateTransformIndex` as public api under the ``xarray.indexes`` namespace. By `Deepak Cherian `_. +- Add convenience methods to :py:class:`~xarray.Coordinates` (:pull:`10318`) + By `Justus Magin `_. Breaking changes From 059512a37dd65f4f7c45e2c53db0f2546fef88f5 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 29 Jun 2025 11:07:53 +0200 Subject: [PATCH 11/12] create docs pages --- doc/api.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/api.rst b/doc/api.rst index df6e87c0cf8..199dec0ea27 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -954,7 +954,11 @@ Coordinates contents Coordinates.to_index Coordinates.assign Coordinates.merge + Coordinates.__or__ Coordinates.copy + Coordinates.drop_vars + Coordinates.rename_dims + Coordinates.rename_vars Comparisons ----------- From 773f1497f10e4ae0d2dc789e00dbcd7d46460c87 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 29 Jun 2025 11:09:00 +0200 Subject: [PATCH 12/12] add back the docstring for `__or__` --- xarray/core/coordinates.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index ccb9293b797..3b2448d4c80 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -564,6 +564,26 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset: 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