Expose "Coordinates" as part of Xarray's public API

doc/api-hidden.rst

@@ -9,17 +9,42 @@
 .. autosummary::
    :toctree: generated/
 
+   Coordinates.from_pandas_multiindex
+   Coordinates.get
+   Coordinates.items
+   Coordinates.keys
+   Coordinates.values
+   Coordinates.dims
+   Coordinates.dtypes
+   Coordinates.variables
+   Coordinates.xindexes
+   Coordinates.indexes
+   Coordinates.to_dataset
+   Coordinates.to_index
+   Coordinates.update
+   Coordinates.merge
+   Coordinates.merge_coords
+   Coordinates.copy
+   Coordinates.equals
+   Coordinates.identical
+
    core.coordinates.DatasetCoordinates.get
    core.coordinates.DatasetCoordinates.items
    core.coordinates.DatasetCoordinates.keys
-   core.coordinates.DatasetCoordinates.merge
-   core.coordinates.DatasetCoordinates.to_dataset
-   core.coordinates.DatasetCoordinates.to_index
-   core.coordinates.DatasetCoordinates.update
    core.coordinates.DatasetCoordinates.values
    core.coordinates.DatasetCoordinates.dims
-   core.coordinates.DatasetCoordinates.indexes
+   core.coordinates.DatasetCoordinates.dtypes
    core.coordinates.DatasetCoordinates.variables
+   core.coordinates.DatasetCoordinates.xindexes
+   core.coordinates.DatasetCoordinates.indexes
+   core.coordinates.DatasetCoordinates.to_dataset
+   core.coordinates.DatasetCoordinates.to_index
+   core.coordinates.DatasetCoordinates.update
+   core.coordinates.DatasetCoordinates.merge
+   core.coordinates.DatasetCoordinates.merge_coords
+   core.coordinates.DataArrayCoordinates.copy
+   core.coordinates.DatasetCoordinates.equals
+   core.coordinates.DatasetCoordinates.identical
 
    core.rolling.DatasetCoarsen.boundary
    core.rolling.DatasetCoarsen.coord_func
@@ -47,14 +72,20 @@
    core.coordinates.DataArrayCoordinates.get
    core.coordinates.DataArrayCoordinates.items
    core.coordinates.DataArrayCoordinates.keys
-   core.coordinates.DataArrayCoordinates.merge
-   core.coordinates.DataArrayCoordinates.to_dataset
-   core.coordinates.DataArrayCoordinates.to_index
-   core.coordinates.DataArrayCoordinates.update
    core.coordinates.DataArrayCoordinates.values
    core.coordinates.DataArrayCoordinates.dims
-   core.coordinates.DataArrayCoordinates.indexes
+   core.coordinates.DataArrayCoordinates.dtypes
    core.coordinates.DataArrayCoordinates.variables
+   core.coordinates.DataArrayCoordinates.xindexes
+   core.coordinates.DataArrayCoordinates.indexes
+   core.coordinates.DataArrayCoordinates.to_dataset
+   core.coordinates.DataArrayCoordinates.to_index
+   core.coordinates.DataArrayCoordinates.update
+   core.coordinates.DataArrayCoordinates.merge
+   core.coordinates.DataArrayCoordinates.merge_coords
+   core.coordinates.DataArrayCoordinates.copy
+   core.coordinates.DataArrayCoordinates.equals
+   core.coordinates.DataArrayCoordinates.identical
 
    core.rolling.DataArrayCoarsen.boundary
    core.rolling.DataArrayCoarsen.coord_func

doc/api.rst

@@ -1085,6 +1085,7 @@ Advanced API
 .. autosummary::
    :toctree: generated/
 
+   Coordinates
    Dataset.variables
    DataArray.variable
    Variable

doc/whats-new.rst

@@ -22,6 +22,20 @@ v2023.07.1 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- :py:class:`Coordinates` can now be constructed independently of any Dataset or
+  DataArray (it is also returned by the :py:attr:`Dataset.coords` and
+  :py:attr:`DataArray.coords` properties). ``Coordinates`` objects are useful for
+  passing both coordinate variables and indexes to new Dataset / DataArray objects,
+  e.g., via their constructor or via :py:meth:`Dataset.assign_coords`. We may also
+  wrap coordinate variables in a ``Coordinates`` object in order to skip
+  the automatic creation of (pandas) indexes for dimension coordinates.
+  The :py:class:`Coordinates.from_pandas_multiindex` constructor may be used to
+  create coordinates directly from a :py:class:`pandas.MultiIndex` object (it is
+  preferred over passing it directly as coordinate data, which may be deprecated soon).
+  Like Dataset and DataArray objects, ``Coordinates`` objects may now be used in
+  :py:func:`align` and :py:func:`merge`.
+  (:issue:`6392`, :pull:`7368`).
+  By `Benoît Bovy <https://github.com/benbovy>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/__init__.py

@@ -26,6 +26,7 @@
     where,
 )
 from xarray.core.concat import concat
+from xarray.core.coordinates import Coordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
 from xarray.core.extensions import (
@@ -35,7 +36,7 @@
 from xarray.core.merge import Context, MergeError, merge
 from xarray.core.options import get_options, set_options
 from xarray.core.parallel import map_blocks
-from xarray.core.variable import Coordinate, IndexVariable, Variable, as_variable
+from xarray.core.variable import IndexVariable, Variable, as_variable
 from xarray.util.print_versions import show_versions
 
 try:
@@ -98,6 +99,7 @@
     "CFTimeIndex",
     "Context",
     "Coordinate",
+    "Coordinates",
     "DataArray",
     "Dataset",
     "IndexVariable",

xarray/core/alignment.py

@@ -5,13 +5,12 @@
 from collections import defaultdict
 from collections.abc import Hashable, Iterable, Mapping
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, cast
 
 import numpy as np
 import pandas as pd
 
 from xarray.core import dtypes
-from xarray.core.common import DataWithCoords
 from xarray.core.indexes import (
     Index,
     Indexes,
@@ -20,15 +19,14 @@
     indexes_all_equal,
     safe_cast_to_index,
 )
+from xarray.core.types import T_Alignable
 from xarray.core.utils import is_dict_like, is_full_slice
 from xarray.core.variable import Variable, as_compatible_data, calculate_dimensions
 
 if TYPE_CHECKING:
     from xarray.core.dataarray import DataArray
     from xarray.core.dataset import Dataset
-    from xarray.core.types import JoinOptions, T_DataArray, T_Dataset, T_DataWithCoords
-
-DataAlignable = TypeVar("DataAlignable", bound=DataWithCoords)
+    from xarray.core.types import JoinOptions, T_DataArray, T_Dataset
 
 
 def reindex_variables(
@@ -92,7 +90,7 @@ def reindex_variables(
 NormalizedIndexVars = dict[MatchingIndexKey, dict[Hashable, Variable]]
 
 
-class Aligner(Generic[DataAlignable]):
+class Aligner(Generic[T_Alignable]):
     """Implements all the complex logic for the re-indexing and alignment of Xarray
     objects.
 
@@ -105,8 +103,8 @@ class Aligner(Generic[DataAlignable]):
 
     """
 
-    objects: tuple[DataAlignable, ...]
-    results: tuple[DataAlignable, ...]
+    objects: tuple[T_Alignable, ...]
+    results: tuple[T_Alignable, ...]
     objects_matching_indexes: tuple[dict[MatchingIndexKey, Index], ...]
     join: str
     exclude_dims: frozenset[Hashable]
@@ -127,7 +125,7 @@ class Aligner(Generic[DataAlignable]):
 
     def __init__(
         self,
-        objects: Iterable[DataAlignable],
+        objects: Iterable[T_Alignable],
         join: str = "inner",
         indexes: Mapping[Any, Any] | None = None,
         exclude_dims: Iterable = frozenset(),
@@ -510,7 +508,7 @@ def _get_dim_pos_indexers(
 
     def _get_indexes_and_vars(
         self,
-        obj: DataAlignable,
+        obj: T_Alignable,
         matching_indexes: dict[MatchingIndexKey, Index],
     ) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]:
         new_indexes = {}
@@ -533,13 +531,13 @@ def _get_indexes_and_vars(
 
     def _reindex_one(
         self,
-        obj: DataAlignable,
+        obj: T_Alignable,
         matching_indexes: dict[MatchingIndexKey, Index],
-    ) -> DataAlignable:
+    ) -> T_Alignable:
         new_indexes, new_variables = self._get_indexes_and_vars(obj, matching_indexes)
         dim_pos_indexers = self._get_dim_pos_indexers(matching_indexes)
 
-        new_obj = obj._reindex_callback(
+        return obj._reindex_callback(
             self,
             dim_pos_indexers,
             new_variables,
@@ -548,8 +546,6 @@ def _reindex_one(
             self.exclude_dims,
             self.exclude_vars,
         )
-        new_obj.encoding = obj.encoding
-        return new_obj
 
     def reindex_all(self) -> None:
         self.results = tuple(
@@ -581,13 +577,13 @@ def align(self) -> None:
 
 
 def align(
-    *objects: DataAlignable,
+    *objects: T_Alignable,
     join: JoinOptions = "inner",
     copy: bool = True,
     indexes=None,
     exclude=frozenset(),
     fill_value=dtypes.NA,
-) -> tuple[DataAlignable, ...]:
+) -> tuple[T_Alignable, ...]:
     """
     Given any number of Dataset and/or DataArray objects, returns new
     objects with aligned indexes and dimension sizes.
@@ -801,14 +797,15 @@ def deep_align(
 
     This function is not public API.
     """
+    from xarray.core.coordinates import Coordinates
     from xarray.core.dataarray import DataArray
     from xarray.core.dataset import Dataset
 
     if indexes is None:
         indexes = {}
 
     def is_alignable(obj):
-        return isinstance(obj, (DataArray, Dataset))
+        return isinstance(obj, (Coordinates, DataArray, Dataset))
 
     positions = []
     keys = []
@@ -866,15 +863,15 @@ def is_alignable(obj):
 
 
 def reindex(
-    obj: DataAlignable,
+    obj: T_Alignable,
     indexers: Mapping[Any, Any],
     method: str | None = None,
     tolerance: int | float | Iterable[int | float] | None = None,
     copy: bool = True,
     fill_value: Any = dtypes.NA,
     sparse: bool = False,
     exclude_vars: Iterable[Hashable] = frozenset(),
-) -> DataAlignable:
+) -> T_Alignable:
     """Re-index either a Dataset or a DataArray.
 
     Not public API.
@@ -905,13 +902,13 @@ def reindex(
 
 
 def reindex_like(
-    obj: DataAlignable,
+    obj: T_Alignable,
     other: Dataset | DataArray,
     method: str | None = None,
     tolerance: int | float | Iterable[int | float] | None = None,
     copy: bool = True,
     fill_value: Any = dtypes.NA,
-) -> DataAlignable:
+) -> T_Alignable:
     """Re-index either a Dataset or a DataArray like another Dataset/DataArray.
 
     Not public API.
@@ -953,8 +950,8 @@ def _get_broadcast_dims_map_common_coords(args, exclude):
 
 
 def _broadcast_helper(
-    arg: T_DataWithCoords, exclude, dims_map, common_coords
-) -> T_DataWithCoords:
+    arg: T_Alignable, exclude, dims_map, common_coords
+) -> T_Alignable:
     from xarray.core.dataarray import DataArray
     from xarray.core.dataset import Dataset
 
@@ -984,16 +981,16 @@ def _broadcast_dataset(ds: T_Dataset) -> T_Dataset:
 
     # remove casts once https://github.com/python/mypy/issues/12800 is resolved
     if isinstance(arg, DataArray):
-        return cast("T_DataWithCoords", _broadcast_array(arg))
+        return cast(T_Alignable, _broadcast_array(arg))
     elif isinstance(arg, Dataset):
-        return cast("T_DataWithCoords", _broadcast_dataset(arg))
+        return cast(T_Alignable, _broadcast_dataset(arg))
     else:
         raise ValueError("all input must be Dataset or DataArray objects")
 
 
 # TODO: this typing is too restrictive since it cannot deal with mixed
 # DataArray and Dataset types...? Is this a problem?
-def broadcast(*args: T_DataWithCoords, exclude=None) -> tuple[T_DataWithCoords, ...]:
+def broadcast(*args: T_Alignable, exclude=None) -> tuple[T_Alignable, ...]:
     """Explicitly broadcast any number of DataArray or Dataset objects against
     one another.
 

xarray/core/common.py

@@ -14,7 +14,6 @@
 from xarray.core.indexing import BasicIndexer, ExplicitlyIndexed
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.parallelcompat import get_chunked_array_type, guess_chunkmanager
-from xarray.core.pdcompat import _convert_base_to_offset
 from xarray.core.pycompat import is_chunked_array
 from xarray.core.utils import (
     Frozen,
@@ -609,9 +608,17 @@ def assign_coords(
         Dataset.swap_dims
         Dataset.set_coords
         """
+        from xarray.core.coordinates import Coordinates
+
         coords_combined = either_dict_or_kwargs(coords, coords_kwargs, "assign_coords")
         data = self.copy(deep=False)
-        results: dict[Hashable, Any] = self._calc_assign_results(coords_combined)
+
+        results: Coordinates | dict[Hashable, Any]
+        if isinstance(coords, Coordinates):
+            results = coords
+        else:
+            results = self._calc_assign_results(coords_combined)
+
         data.coords.update(results)
         return data
 
@@ -952,6 +959,7 @@ def _resample(
 
         from xarray.core.dataarray import DataArray
         from xarray.core.groupby import ResolvedTimeResampleGrouper, TimeResampleGrouper
+        from xarray.core.pdcompat import _convert_base_to_offset
         from xarray.core.resample import RESAMPLE_DIM
 
         if keep_attrs is not None: