lint, public API and docstrings

pydata · dcherian · Feb 14, 2025 · Sep 24, 2024 · Sep 24, 2024 · Sep 24, 2024
commit 0b545cf61cf192cd1037e2e1d312921a8ab5843c
diff --git a/xarray/__init__.py b/xarray/__init__.py
@@ -30,6 +30,7 @@
     where,
 )
 from xarray.core.concat import concat
+from xarray.core.coordinate_transform import CoordinateTransform
 from xarray.core.coordinates import Coordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
@@ -109,6 +110,7 @@
     "CFTimeIndex",
     "Context",
     "Coordinates",
+    "CoordinateTransform",
     "DataArray",
     "Dataset",
     "DataTree",

diff --git a/xarray/core/coordinate_transform.py b/xarray/core/coordinate_transform.py
@@ -1,4 +1,5 @@
-from typing import Any, Iterable, Hashable, Mapping
+from collections.abc import Hashable, Iterable, Mapping
+from typing import Any
 
 import numpy as np
 
@@ -15,11 +16,14 @@ def __init__(
         self,
         coord_names: Iterable[Hashable],
         dim_size: Mapping[str, int],
-        dtype: Any = np.dtype(np.float64),
+        dtype: Any = None,
     ):
         self.coord_names = tuple(coord_names)
         self.dims = tuple(dim_size)
         self.dim_size = dict(dim_size)
+
+        if dtype is None:
+            dtype = np.dtype(np.float64)
         self.dtype = dtype
 
     def forward(self, dim_positions: dict[str, Any]) -> dict[Hashable, Any]:
@@ -61,7 +65,7 @@ def equals(self, other: "CoordinateTransform") -> bool:
         raise NotImplementedError
 
     def generate_coords(self, dims: tuple[str] | None = None) -> dict[Hashable, Any]:
-        """Returns all "world" coordinate labels."""
+        """Compute all coordinate labels at once."""
         if dims is None:
             dims = self.dims
 

diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py
@@ -12,6 +12,7 @@
 from xarray.core import formatting, nputils, utils
 from xarray.core.coordinate_transform import CoordinateTransform
 from xarray.core.indexing import (
+    CoordinateTransformIndexingAdapter,
     IndexSelResult,
     PandasIndexingAdapter,
     PandasMultiIndexingAdapter,
@@ -25,6 +26,7 @@
 )
 
 if TYPE_CHECKING:
+    from xarray.core.coordinate import Coordinates
     from xarray.core.types import ErrorOptions, JoinOptions, Self
     from xarray.core.variable import Variable
 
@@ -1374,8 +1376,13 @@ def rename(self, name_dict, dims_dict):
 
 
 class CoordinateTransformIndex(Index):
-    """Xarray index abstract class for transformation between "pixel"
-    and "world" coordinates.
+    """Helper class for creating Xarray indexes based on coordinate transforms.
+
+    - wraps a :py:class:`CoordinateTransform` instance
+    - takes care of creating the index (lazy) coordinates
+    - supports point-wise label-based selection
+    - supports exact alignment only, by comparing indexes based on their transform
+      (not on their explicit coordinate labels)
 
     """
 
@@ -1409,9 +1416,11 @@ def create_variables(
 
     def create_coordinates(self) -> Coordinates:
         # TODO: move this in xarray.Index base class?
+        from xarray.core.coordinates import Coordinates
+
         variables = self.create_variables()
         indexes = {name: self for name in variables}
-        return xr.Coordinates(coords=variables, indexes=indexes)
+        return Coordinates(coords=variables, indexes=indexes)
 
     def isel(
         self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]
@@ -1423,6 +1432,9 @@ def isel(
     def sel(
         self, labels: dict[Any, Any], method=None, tolerance=None
     ) -> IndexSelResult:
+        from xarray.core.dataarray import DataArray
+        from xarray.core.variable import Variable
+
         if method != "nearest":
             raise ValueError(
                 "CoordinateTransformIndex only supports selection with method='nearest'"
@@ -1433,15 +1445,14 @@ def sel(
 
         missing_labels = coord_names_set - labels_set
         if missing_labels:
-            raise ValueError(
-                f"missing labels for coordinate(s): {','.join(missing_labels)}."
-            )
+            missing_labels_str = ",".join([f"{name}" for name in missing_labels])
+            raise ValueError(f"missing labels for coordinate(s): {missing_labels_str}.")
 
         label0_obj = next(iter(labels.values()))
         dim_size0 = getattr(label0_obj, "sizes", None)
 
         is_xr_obj = [
-            isinstance(label, (xr.DataArray, xr.Variable)) for label in labels.values()
+            isinstance(label, DataArray | Variable) for label in labels.values()
         ]
         if not all(is_xr_obj):
             raise TypeError(
@@ -1461,13 +1472,14 @@ def sel(
         dim_positions = self.transform.reverse(coord_labels)
 
         results = {}
+        dims0 = tuple(dim_size0)
         for dim, pos in dim_positions.items():
             if isinstance(label0_obj, Variable):
-                xr_pos = Variable(label.dims, idx)
+                xr_pos = Variable(dims0, pos)
             else:
                 # dataarray
-                xr_pos = DataArray(idx, dims=label.dims)
-            results[dim] = idx
+                xr_pos = DataArray(pos, dims=dims0)
+            results[dim] = xr_pos
 
         return IndexSelResult(results)
 

diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py
@@ -1961,8 +1961,9 @@ def copy(self, deep: bool = True) -> Self:
 
 
 class CoordinateTransformIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
-    """Wrap a CoordinateTransform to support explicit indexing and
-    lazy coordinate labels.
+    """Wrap a CoordinateTransform as a lazy coordinate array.
+
+    Supports explicit indexing (both outer and vectorized).
 
     """
 
@@ -2036,7 +2037,7 @@ def __getitem__(self, indexer: ExplicitIndexer):
         self._check_and_raise_if_non_basic_indexer(indexer)
 
         # also works with basic indexing
-        return self._oindex_get(indexer)
+        return self._oindex_get(OuterIndexer(indexer.tuple))
 
     def __setitem__(self, indexer: ExplicitIndexer, value: Any) -> None:
         raise TypeError(

diff --git a/xarray/indexes/__init__.py b/xarray/indexes/__init__.py
@@ -3,6 +3,11 @@
 
 """
 
-from xarray.core.indexes import Index, PandasIndex, PandasMultiIndex
+from xarray.core.indexes import (
+    CoordinateTransformIndex,
+    Index,
+    PandasIndex,
+    PandasMultiIndex,
+)
 
-__all__ = ["Index", "PandasIndex", "PandasMultiIndex"]
+__all__ = ["CoordinateTransformIndex", "Index", "PandasIndex", "PandasMultiIndex"]