Migrate indexing and broadcasting logic to xarray.namedarray (Part 1)

xarray/core/dataset.py

@@ -119,7 +119,7 @@
     IndexVariable,
     Variable,
     as_variable,
-    broadcast_variables,
+    broadcast_namedarrays,
     calculate_dimensions,
 )
 from xarray.namedarray.parallelcompat import get_chunked_array_type, guess_chunkmanager
@@ -291,7 +291,6 @@ def _maybe_chunk(
     chunked_array_type: str | ChunkManagerEntrypoint | None = None,
     from_array_kwargs=None,
 ):
-
     from xarray.namedarray.daskmanager import DaskManager
 
     if chunks is not None:
@@ -7168,7 +7167,7 @@ def to_dataarray(
         from xarray.core.dataarray import DataArray
 
         data_vars = [self.variables[k] for k in self.data_vars]
-        broadcast_vars = broadcast_variables(*data_vars)
+        broadcast_vars = broadcast_namedarrays(*data_vars)
         data = duck_array_ops.stack([b.data for b in broadcast_vars], axis=0)
 
         dims = (dim,) + broadcast_vars[0].dims

xarray/core/indexing.py

@@ -207,7 +207,7 @@ def map_index_queries(
     return merged
 
 
-def expanded_indexer(key, ndim):
+def expanded_indexer(key, ndim: int) -> tuple[Any, ...]:
     """Given a key for indexing an ndarray, return an equivalent key which is a
     tuple with length equal to the number of dimensions.
 
@@ -882,7 +882,7 @@ def __setitem__(self, indexer: _IndexerKey, value: Any) -> None:
         self.array[indexer] = value
 
 
-def as_indexable(array):
+def as_indexable(array: Any):
     """
     This function always returns a ExplicitlyIndexed subclass,
     so that the vectorized indexing is always possible with the returned

xarray/core/missing.py

@@ -22,7 +22,7 @@
 from xarray.core.options import _get_keep_attrs
 from xarray.core.types import Interp1dOptions, InterpOptions
 from xarray.core.utils import OrderedSet, is_scalar
-from xarray.core.variable import Variable, broadcast_variables
+from xarray.core.variable import Variable, broadcast_namedarrays
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import is_chunked_array
 
@@ -623,7 +623,7 @@ def interp(var, indexes_coords, method: InterpOptions, **kwargs):
         # target dimensions
         dims = list(indexes_coords)
         x, new_x = zip(*[indexes_coords[d] for d in dims])
-        destination = broadcast_variables(*new_x)
+        destination = broadcast_namedarrays(*new_x)
 
         # transpose to make the interpolated axis to the last position
         broadcast_dims = [d for d in var.dims if d not in dims]

xarray/core/utils.py

@@ -56,7 +56,6 @@
     KeysView,
     Mapping,
     MutableMapping,
-    MutableSet,
     ValuesView,
 )
 from enum import Enum
@@ -75,6 +74,7 @@
 import pandas as pd
 
 from xarray.namedarray.utils import (  # noqa: F401
+    OrderedSet,
     ReprObject,
     drop_missing_dims,
     either_dict_or_kwargs,
@@ -521,48 +521,6 @@ def __len__(self) -> int:
         return len(self._keys)
 
 
-class OrderedSet(MutableSet[T]):
-    """A simple ordered set.
-
-    The API matches the builtin set, but it preserves insertion order of elements, like
-    a dict. Note that, unlike in an OrderedDict, equality tests are not order-sensitive.
-    """
-
-    _d: dict[T, None]
-
-    __slots__ = ("_d",)
-
-    def __init__(self, values: Iterable[T] | None = None):
-        self._d = {}
-        if values is not None:
-            self.update(values)
-
-    # Required methods for MutableSet
-
-    def __contains__(self, value: Hashable) -> bool:
-        return value in self._d
-
-    def __iter__(self) -> Iterator[T]:
-        return iter(self._d)
-
-    def __len__(self) -> int:
-        return len(self._d)
-
-    def add(self, value: T) -> None:
-        self._d[value] = None
-
-    def discard(self, value: T) -> None:
-        del self._d[value]
-
-    # Additional methods
-
-    def update(self, values: Iterable[T]) -> None:
-        self._d.update(dict.fromkeys(values))
-
-    def __repr__(self) -> str:
-        return f"{type(self).__name__}({list(self)!r})"
-
-
 class NdimSizeLenMixin:
     """Mixin class that extends a class that defines a ``shape`` property to
     one that also defines ``ndim``, ``size`` and ``__len__``.