Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
eb6dc62
implement `Coordinates` methods modifying coords
keewis May 13, 2025
8a18a52
allow using the binary-or operator (`|`) for merging
keewis May 13, 2025
f8b7f57
tests for `drop_vars`
keewis May 13, 2025
92e1e8a
Merge branch 'main' into coords-methods
keewis May 13, 2025
4626ce5
tests for `rename_dims` and `rename_vars`
keewis May 13, 2025
a370ed9
tests for the merge operator
keewis May 13, 2025
c909ca9
Apply suggestions from code review
keewis Jun 28, 2025
836aeca
Merge branch 'main' into coords-methods
keewis Jun 28, 2025
22c89e3
attempt to fix the typing
keewis Jun 29, 2025
3b9e3e5
make sure we always return a `Coordinates` object
keewis Jun 29, 2025
e53b848
replace docstring by a reference to `Coordinates.merge`
keewis Jun 29, 2025
b204459
Merge branch 'main' into coords-methods
keewis Jun 29, 2025
7dbd26a
changelog
keewis Jun 29, 2025
059512a
create docs pages
keewis Jun 29, 2025
773f149
add back the docstring for `__or__`
keewis Jun 29, 2025
c940f1a
Merge branch 'main' into coords-methods
dcherian Jul 11, 2025
1447ea1
Merge branch 'main' into coords-methods
keewis Jul 19, 2025
7bd7913
add `drop_dims`
keewis Jul 28, 2025
85d685b
typing
keewis Jul 28, 2025
e910466
Merge branch 'main' into coords-methods
keewis Jul 28, 2025
c599d61
undo a bad merge [skip-ci]
keewis Jul 28, 2025
ce177db
document `drop_dims` [skip-ci]
keewis Jul 28, 2025
08624d8
Merge branch 'main' into coords-methods
keewis Aug 2, 2025
4f5f525
Merge branch 'main' into coords-methods
keewis Aug 23, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions doc/api/coordinates.rst
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ and values given by ``DataArray`` objects.
Coordinates.__getitem__
Coordinates.__setitem__
Coordinates.__delitem__
Coordinates.__or__
Coordinates.update
Coordinates.get
Coordinates.items
Expand All @@ -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
-----------
Expand Down
3 changes: 2 additions & 1 deletion doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down
135 changes: 133 additions & 2 deletions xarray/core/coordinates.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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})

Expand Down Expand Up @@ -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).
Expand Down
84 changes: 84 additions & 0 deletions xarray/tests/test_coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading