Add legacy backend indexing tests

Author: dcherian

xarray/backends/common.py

@@ -250,6 +250,12 @@ def robust_getitem(array, key, catch=Exception, max_retries=6, initial_delay=500
 
 
 class BackendArray(NdimSizeLenMixin, indexing.ExplicitlyIndexed):
+    def get_duck_array(self, dtype: np.typing.DTypeLike = None):
+        key = indexing.BasicIndexer((slice(None),) * self.ndim)
+        return self[key]  # type: ignore [index]
+
+
+class NewBackendArray(NdimSizeLenMixin, indexing.ExplicitlyIndexed):
     __slots__ = ("indexing_support",)
 
     def get_duck_array(self, dtype: np.typing.DTypeLike = None):

xarray/backends/netCDF4_.py

@@ -12,8 +12,8 @@
 from xarray import coding
 from xarray.backends.common import (
     BACKEND_ENTRYPOINTS,
-    BackendArray,
     BackendEntrypoint,
+    NewBackendArray,
     WritableCFDataStore,
     _normalize_path,
     datatree_from_dict_with_io_cleanup,
@@ -61,7 +61,7 @@
 NETCDF4_PYTHON_LOCK = combine_locks([NETCDFC_LOCK, HDF5_LOCK])
 
 
-class BaseNetCDF4Array(BackendArray):
+class BaseNetCDF4Array(NewBackendArray):
     __slots__ = ("datastore", "dtype", "shape", "variable_name")
 
     def __init__(self, variable_name, datastore):

xarray/backends/pydap_.py

@@ -8,8 +8,8 @@
 from xarray.backends.common import (
     BACKEND_ENTRYPOINTS,
     AbstractDataStore,
-    BackendArray,
     BackendEntrypoint,
+    NewBackendArray,
     robust_getitem,
 )
 from xarray.backends.store import StoreBackendEntrypoint
@@ -36,7 +36,7 @@
     )
 
 
-class PydapArrayWrapper(BackendArray):
+class PydapArrayWrapper(NewBackendArray):
     indexing_support = indexing.IndexingSupport.BASIC
 
     def __init__(self, array) -> None:

xarray/backends/scipy_.py

@@ -10,8 +10,8 @@
 
 from xarray.backends.common import (
     BACKEND_ENTRYPOINTS,
-    BackendArray,
     BackendEntrypoint,
+    NewBackendArray,
     WritableCFDataStore,
     _normalize_path,
 )
@@ -59,7 +59,7 @@ def _decode_attrs(d):
     return {k: v if k == "_FillValue" else _decode_string(v) for (k, v) in d.items()}
 
 
-class ScipyArrayWrapper(BackendArray):
+class ScipyArrayWrapper(NewBackendArray):
     indexing_support = indexing.IndexingSupport.OUTER_1VECTOR
 
     def __init__(self, variable_name, datastore):

xarray/backends/zarr.py

@@ -14,8 +14,8 @@
 from xarray.backends.common import (
     BACKEND_ENTRYPOINTS,
     AbstractWritableDataStore,
-    BackendArray,
     BackendEntrypoint,
+    NewBackendArray,
     _encode_variable_name,
     _normalize_path,
     datatree_from_dict_with_io_cleanup,
@@ -185,7 +185,7 @@ def encode_zarr_attr_value(value):
     return encoded
 
 
-class ZarrArrayWrapper(BackendArray):
+class ZarrArrayWrapper(NewBackendArray):
     indexing_support = indexing.IndexingSupport.VECTORIZED
 
     def __init__(self, zarr_array):

xarray/core/indexing.py

@@ -3,7 +3,6 @@
 import enum
 import functools
 import operator
-import warnings
 from collections import Counter, defaultdict
 from collections.abc import Callable, Hashable, Iterable, Mapping
 from contextlib import suppress
@@ -23,6 +22,7 @@
 from xarray.core.utils import (
     NDArrayMixin,
     either_dict_or_kwargs,
+    emit_user_level_warning,
     get_valid_numpy_dtype,
     is_duck_array,
     is_duck_dask_array,
@@ -609,7 +609,7 @@ def __getitem__(self, key) -> Any:
 BackendArray_fallback_warning_message = (
     "The array `{0}` does not support indexing using the .vindex and .oindex properties. "
     "The __getitem__ method is being used instead. This fallback behavior will be "
-    "removed in a future version. Please ensure that the backend array `{1}` implements "
+    "removed in a future version. Please ensure that the backend array `{0}` implements "
     "support for the .vindex and .oindex properties to avoid potential issues."
 )
 
@@ -671,21 +671,8 @@ def shape(self) -> _Shape:
         return self._shape
 
     def get_duck_array(self) -> Any:
-        try:
-            array = apply_indexer(self.array, self.key)
-        except NotImplementedError as _:
-            # If the array is not an ExplicitlyIndexedNDArrayMixin,
-            # it may wrap a BackendArray subclass that doesn't implement .oindex and .vindex. so use its __getitem__
-            warnings.warn(
-                BackendArray_fallback_warning_message.format(
-                    self.array.__class__.__name__, self.array.__class__.__name__
-                ),
-                category=DeprecationWarning,
-                stacklevel=2,
-            )
-            array = self.array[self.key]
-
-        # self.array[self.key] is now a numpy array when
+        array = apply_indexer(self.array, self.key)
+        # array[self.key] is now a numpy array when
         # self.array is a BackendArray subclass
         # and self.key is BasicIndexer((slice(None, None, None),))
         # so we need the explicit check for ExplicitlyIndexed
@@ -752,21 +739,9 @@ def shape(self) -> _Shape:
         return np.broadcast(*self.key.tuple).shape
 
     def get_duck_array(self) -> Any:
-        try:
-            array = apply_indexer(self.array, self.key)
-        except NotImplementedError as _:
-            # If the array is not an ExplicitlyIndexedNDArrayMixin,
-            # it may wrap a BackendArray subclass that doesn't implement .oindex and .vindex. so use its __getitem__
-            warnings.warn(
-                BackendArray_fallback_warning_message.format(
-                    self.array.__class__.__name__, self.array.__class__.__name__
-                ),
-                category=PendingDeprecationWarning,
-                stacklevel=2,
-            )
-            array = self.array[self.key]
+        array = apply_indexer(self.array, self.key)
 
-        # self.array[self.key] is now a numpy array when
+        # array is now a numpy array when
         # self.array is a BackendArray subclass
         # and self.key is BasicIndexer((slice(None, None, None),))
         # so we need the explicit check for ExplicitlyIndexed
@@ -1136,6 +1111,7 @@ def vectorized_indexing_adapter(
     )
 
 
+# TODO: deprecate and delete this method once it is no longer used externally
 def explicit_indexing_adapter(
     key: ExplicitIndexer,
     shape: _Shape,
@@ -1163,26 +1139,36 @@ def explicit_indexing_adapter(
     -------
     Indexing result, in the form of a duck numpy-array.
     """
-    # TODO: raise PendingDeprecationWarning here.
-    if isinstance(key, VectorizedIndexer):
-        return vectorized_indexing_adapter(
-            key.tuple, shape, indexing_support, raw_indexing_method
-        )
-    elif isinstance(key, OuterIndexer):
-        return outer_indexing_adapter(
-            key.tuple, shape, indexing_support, raw_indexing_method
-        )
-    elif isinstance(key, BasicIndexer):
-        return basic_indexing_adapter(
-            key.tuple, shape, indexing_support, raw_indexing_method
-        )
-    raise TypeError(f"unexpected key type: {key}")
+
+    # If the array is not an ExplicitlyIndexedNDArrayMixin,
+    # it may wrap a BackendArray subclass that doesn't implement .oindex and .vindex. so use its __getitem__
+    emit_user_level_warning(
+        BackendArray_fallback_warning_message.format(""),
+        category=PendingDeprecationWarning,
+    )
+    raw_key, numpy_indices = decompose_indexer(key, shape, indexing_support)
+    result = raw_indexing_method(raw_key.tuple)
+    if numpy_indices.tuple:
+        indexable = NumpyIndexingAdapter(result)
+        result = apply_indexer(indexable, numpy_indices)
+    return result
 
 
 def apply_indexer(
     indexable: ExplicitlyIndexedNDArrayMixin, indexer: ExplicitIndexer
 ) -> Any:
     """Apply an indexer to an indexable object."""
+    if not hasattr(indexable, "vindex") and not hasattr(indexable, "oindex"):
+        # This path is used by Lazily*IndexedArray.get_duck_array()
+        classname = type(indexable).__name__
+        # If the array is not an ExplicitlyIndexedNDArrayMixin,
+        # it may wrap a BackendArray subclass that doesn't implement .oindex and .vindex. so use its __getitem__
+        emit_user_level_warning(
+            BackendArray_fallback_warning_message.format(classname),
+            category=PendingDeprecationWarning,
+        )
+        return indexable[indexer]
+
     if isinstance(indexer, VectorizedIndexer):
         return indexable.vindex[indexer.tuple]
     elif isinstance(indexer, OuterIndexer):
@@ -1206,6 +1192,7 @@ def set_with_indexer(indexable, indexer: ExplicitIndexer, value: Any) -> None:
         indexable[indexer.tuple] = value
 
 
+# TODO: delete this method once explicit_indexing_adapter is no longer used externally
 def decompose_indexer(
     indexer: ExplicitIndexer, shape: _Shape, indexing_support: IndexingSupport
 ) -> tuple[ExplicitIndexer, ExplicitIndexer]:

xarray/tests/test_backends.py

@@ -39,6 +39,7 @@
     open_mfdataset,
     save_mfdataset,
 )
+from xarray.backends.common import BackendArray as LegacyBackendArray
 from xarray.backends.common import robust_getitem
 from xarray.backends.h5netcdf_ import H5netcdfBackendEntrypoint
 from xarray.backends.netcdf3 import _nc3_dtype_coercions
@@ -53,6 +54,7 @@
 from xarray.coding.variables import SerializationWarning
 from xarray.conventions import encode_dataset_coordinates
 from xarray.core import indexing
+from xarray.core.indexing import IndexingSupport
 from xarray.core.options import set_options
 from xarray.core.utils import module_available
 from xarray.namedarray.pycompat import array_type
@@ -352,6 +354,9 @@ def test_dtype_coercion_error(self) -> None:
 
 
 class BackendIndexingTestsMixin:
+    def roundtrip(self, ds: Dataset, open_kwargs=None) -> Dataset:
+        raise NotImplementedError
+
     def test_orthogonal_indexing(self) -> None:
         in_memory = create_test_data()
         with self.roundtrip(in_memory) as on_disk:
@@ -6491,3 +6496,83 @@ def test_zarr_safe_chunk_region(tmp_path):
     chunk = ds.isel(region)
     chunk = chunk.chunk()
     chunk.chunk().to_zarr(store, region=region)
+
+
+class LegacyBackendArrayWrapper(LegacyBackendArray):
+    def __init__(self, array: np.ndarray, indexing_support: IndexingSupport):
+        self.shape = array.shape
+        self.dtype = array.dtype
+        self.array = array
+        self.indexing_support = indexing_support
+
+    def __getitem__(self, key: indexing.ExplicitIndexer):
+        return indexing.explicit_indexing_adapter(
+            key, self.shape, self.indexing_support, self._getitem
+        )
+
+    def _getitem(self, key: tuple[Any, ...]) -> np.ndarray:
+        return self.array[key]
+
+
+def indexing_tests(*, indexing_support: IndexingSupport):
+    def wrapper(cls):
+        class NewClass(cls):
+            cls.indexing_support = indexing_support
+
+            def roundtrip(self, ds: Dataset, *, open_kwargs=None) -> Dataset:
+                ds = ds.copy(deep=True)
+                for name in list(ds.data_vars) + list(
+                    set(ds.coords) - set(ds.xindexes)
+                ):
+                    var = ds._variables[name]
+                    ds._variables[name] = var.copy(
+                        # These tests assume that indexing is lazy (checks ._in_memory),
+                        # so wrapping by LazilyIndexedArray is required.
+                        data=indexing.LazilyIndexedArray(
+                            LegacyBackendArrayWrapper(var.data, self.indexing_support)
+                        )
+                    )
+                return ds
+
+            def test_vectorized_indexing_negative_step(self) -> None:
+                with pytest.warns(PendingDeprecationWarning):
+                    super().test_vectorized_indexing_negative_step()
+
+            def test_isel_dataarray(self) -> None:
+                with pytest.warns(PendingDeprecationWarning):
+                    super().test_isel_dataarray()
+
+            def test_vectorized_indexing(self) -> None:
+                with pytest.warns(PendingDeprecationWarning):
+                    super().test_vectorized_indexing()
+
+            def test_orthogonal_indexing(self) -> None:
+                with pytest.warns(PendingDeprecationWarning):
+                    super().test_orthogonal_indexing()
+
+            def test_outer_indexing_reversed(self) -> None:
+                with pytest.warns(PendingDeprecationWarning):
+                    super().test_outer_indexing_reversed()
+
+        return NewClass
+
+    return wrapper
+
+
+@indexing_tests(indexing_support=IndexingSupport.BASIC)
+class TestBasicIndexingLegacyBackend(BackendIndexingTestsMixin):
+    pass
+
+
+@indexing_tests(indexing_support=IndexingSupport.OUTER_1VECTOR)
+class TestOuter1VectorIndexingLegacyBackend(BackendIndexingTestsMixin):
+    pass
+
+
+# @indexing_tests(indexing_support=IndexingSupport.OUTER)
+# class TestOuterIndexingLegacyBackend(BackendIndexingTestsMixin):
+#     pass
+
+# @indexing_tests(indexing_support=IndexingSupport.VECTORIZED)
+# class TestVectorizedIndexingLegacyBackend(BackendIndexingTestsMixin):
+#     pass