Use tuples for indexing

Author: dcherian

xarray/backends/common.py

@@ -22,6 +22,7 @@
 
     from xarray.core.dataset import Dataset
     from xarray.core.types import NestedSequence
+    from xarray.namedarray._typing import _IndexerKey
 
 # Create a logger object, but don't add any handlers. Leave that to user code.
 logger = logging.getLogger(__name__)
@@ -219,18 +220,18 @@ def robust_getitem(array, key, catch=Exception, max_retries=6, initial_delay=500
 
 
 class BackendArray(NdimSizeLenMixin, indexing.ExplicitlyIndexed):
-    __slots__ = ()
+    __slots__ = ("indexing_support",)
 
     def get_duck_array(self, dtype: np.typing.DTypeLike = None):
-        key = indexing.BasicIndexer((slice(None),) * self.ndim)
+        key = (slice(None),) * self.ndim
         return self[key]  # type: ignore [index]
 
-    def _oindex_get(self, key: indexing.OuterIndexer):
+    def _oindex_get(self, key: _IndexerKey) -> Any:
         raise NotImplementedError(
             f"{self.__class__.__name__}._oindex_get method should be overridden"
         )
 
-    def _vindex_get(self, key: indexing.VectorizedIndexer):
+    def _vindex_get(self, key: _IndexerKey) -> Any:
         raise NotImplementedError(
             f"{self.__class__.__name__}._vindex_get method should be overridden"
         )

xarray/backends/h5netcdf_.py

@@ -44,26 +44,33 @@
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
     from xarray.core.datatree import DataTree
+    from xarray.namedarray._typing import (
+        _BasicIndexerKey,
+        _OuterIndexerKey,
+        _VectorizedIndexerKey,
+    )
 
 
 class H5NetCDFArrayWrapper(BaseNetCDF4Array):
+    indexing_support = indexing.IndexingSupport.OUTER_1VECTOR
+
     def get_array(self, needs_lock=True):
         ds = self.datastore._acquire(needs_lock)
         return ds.variables[self.variable_name]
 
-    def _oindex_get(self, key: indexing.OuterIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
+    def _oindex_get(self, key: _OuterIndexerKey) -> Any:
+        return indexing.outer_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
-    def _vindex_get(self, key: indexing.VectorizedIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
+    def _vindex_get(self, key: _VectorizedIndexerKey) -> Any:
+        return indexing.vectorized_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
-    def __getitem__(self, key: indexing.BasicIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
+    def __getitem__(self, key: _BasicIndexerKey) -> Any:
+        return indexing.basic_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
     def _getitem(self, key):

xarray/backends/netCDF4_.py

@@ -49,6 +49,11 @@
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
     from xarray.core.datatree import DataTree
+    from xarray.namedarray._typing import (
+        _BasicIndexerKey,
+        _OuterIndexerKey,
+        _VectorizedIndexerKey,
+    )
 
 # This lookup table maps from dtype.byteorder to a readable endian
 # string used by netCDF4.
@@ -89,7 +94,7 @@ def get_array(self, needs_lock=True):
 
 
 class NetCDF4ArrayWrapper(BaseNetCDF4Array):
-    __slots__ = ()
+    indexing_support = indexing.IndexingSupport.OUTER
 
     def get_array(self, needs_lock=True):
         ds = self.datastore._acquire(needs_lock)
@@ -100,19 +105,19 @@ def get_array(self, needs_lock=True):
             variable.set_auto_chartostring(False)
         return variable
 
-    def _oindex_get(self, key: indexing.OuterIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER, self._getitem
+    def _oindex_get(self, key: _OuterIndexerKey):
+        return indexing.outer_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
-    def _vindex_get(self, key: indexing.VectorizedIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER, self._getitem
+    def _vindex_get(self, key: _VectorizedIndexerKey):
+        return indexing.vectorized_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
-    def __getitem__(self, key: indexing.BasicIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER, self._getitem
+    def __getitem__(self, key: _BasicIndexerKey):
+        return indexing.basic_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
     def _getitem(self, key):

xarray/backends/pydap_.py

@@ -29,10 +29,17 @@
     from io import BufferedIOBase
 
     from xarray.core.dataset import Dataset
+    from xarray.namedarray._typing import (
+        _BasicIndexerKey,
+        _OuterIndexerKey,
+        _VectorizedIndexerKey,
+    )
 
 
 class PydapArrayWrapper(BackendArray):
-    def __init__(self, array):
+    indexing_support = indexing.IndexingSupport.BASIC
+
+    def __init__(self, array) -> None:
         self.array = array
 
     @property
@@ -43,19 +50,19 @@ def shape(self) -> tuple[int, ...]:
     def dtype(self):
         return self.array.dtype
 
-    def _oindex_get(self, key: indexing.OuterIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
+    def _oindex_get(self, key: _OuterIndexerKey) -> Any:
+        return indexing.outer_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
-    def _vindex_get(self, key: indexing.VectorizedIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
+    def _vindex_get(self, key: _VectorizedIndexerKey) -> Any:
+        return indexing.vectorized_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
-    def __getitem__(self, key: indexing.BasicIndexer):
-        return indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
+    def __getitem__(self, key: _BasicIndexerKey) -> Any:
+        return indexing.basic_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
 
     def _getitem(self, key):

xarray/backends/scipy_.py

@@ -38,6 +38,11 @@
 
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
+    from xarray.namedarray._typing import (
+        _BasicIndexerKey,
+        _OuterIndexerKey,
+        _VectorizedIndexerKey,
+    )
 
 
 HAS_NUMPY_2_0 = module_available("numpy", minversion="2.0.0.dev0")
@@ -56,6 +61,8 @@ def _decode_attrs(d):
 
 
 class ScipyArrayWrapper(BackendArray):
+    indexing_support = indexing.IndexingSupport.OUTER_1VECTOR
+
     def __init__(self, variable_name, datastore):
         self.datastore = datastore
         self.variable_name = variable_name
@@ -85,25 +92,25 @@ def _getitem(self, key):
             data = self.get_variable(needs_lock=False).data
             return data[key]
 
-    def _vindex_get(self, key: indexing.VectorizedIndexer):
-        data = indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
+    def _vindex_get(self, key: _VectorizedIndexerKey) -> Any:
+        data = indexing.vectorized_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
         return self._finalize_result(data)
 
-    def _oindex_get(self, key: indexing.OuterIndexer):
-        data = indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
+    def _oindex_get(self, key: _OuterIndexerKey) -> Any:
+        data = indexing.outer_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
         return self._finalize_result(data)
 
-    def __getitem__(self, key):
-        data = indexing.explicit_indexing_adapter(
-            key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
+    def __getitem__(self, key: _BasicIndexerKey) -> Any:
+        data = indexing.basic_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
         )
         return self._finalize_result(data)
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key, value) -> None:
         with self.datastore.lock:
             data = self.get_variable(needs_lock=False)
             try:

xarray/backends/zarr.py

@@ -43,6 +43,11 @@
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
     from xarray.core.datatree import DataTree
+    from xarray.namedarray._typing import (
+        _BasicIndexerKey,
+        _OuterIndexerKey,
+        _VectorizedIndexerKey,
+    )
 
 
 def _get_mappers(*, storage_options, store, chunk_store):
@@ -182,7 +187,7 @@ def encode_zarr_attr_value(value):
 
 
 class ZarrArrayWrapper(BackendArray):
-    __slots__ = ("_array", "dtype", "shape")
+    indexing_support = indexing.IndexingSupport.VECTORIZED
 
     def __init__(self, zarr_array):
         # some callers attempt to evaluate an array if an `array` property exists on the object.
@@ -205,37 +210,28 @@ def __init__(self, zarr_array):
     def get_array(self):
         return self._array
 
-    def _oindex_get(self, key: indexing.OuterIndexer):
+    def _oindex_get(self, key: _OuterIndexerKey) -> Any:
         def raw_indexing_method(key):
             return self._array.oindex[key]
 
-        return indexing.explicit_indexing_adapter(
-            key,
-            self._array.shape,
-            indexing.IndexingSupport.VECTORIZED,
-            raw_indexing_method,
+        return indexing.outer_indexing_adapter(
+            key, self._array.shape, self.indexing_support, raw_indexing_method
         )
 
-    def _vindex_get(self, key: indexing.VectorizedIndexer):
+    def _vindex_get(self, key: _VectorizedIndexerKey) -> Any:
         def raw_indexing_method(key):
             return self._array.vindex[key]
 
-        return indexing.explicit_indexing_adapter(
-            key,
-            self._array.shape,
-            indexing.IndexingSupport.VECTORIZED,
-            raw_indexing_method,
+        return indexing.vectorized_indexing_adapter(
+            key, self._array.shape, self.indexing_support, raw_indexing_method
         )
 
-    def __getitem__(self, key: indexing.BasicIndexer):
+    def __getitem__(self, key: _BasicIndexerKey) -> Any:
         def raw_indexing_method(key):
             return self._array[key]
 
-        return indexing.explicit_indexing_adapter(
-            key,
-            self._array.shape,
-            indexing.IndexingSupport.VECTORIZED,
-            raw_indexing_method,
+        return indexing.basic_indexing_adapter(
+            key, self._array.shape, self.indexing_support, raw_indexing_method
         )
 
         # if self.ndim == 0:

xarray/coding/strings.py

@@ -256,10 +256,8 @@ def _oindex_get(self, key: indexing.OuterIndexer):
         return _numpy_char_to_bytes(self.array.oindex[key])
 
     def __getitem__(self, key: _IndexerKey):
-        from xarray.core.indexing import BasicIndexer
-
         # require slicing the last dimension completely
         indexer = indexing.expanded_indexer(key, self.array.ndim)
         if indexer[-1] != slice(None):
             raise IndexError("too many indices")
-        return _numpy_char_to_bytes(self.array[BasicIndexer(indexer)])
+        return _numpy_char_to_bytes(self.array[indexer])