check that vectorized indexing also works

Author: keewis

xarray_array_testing/indexing.py

@@ -2,6 +2,7 @@
 
 import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
+import xarray as xr
 import xarray.testing.strategies as xrst
 from hypothesis import given
 
@@ -51,6 +52,24 @@ def orthogonal_indexers(draw, sizes, indexer_types):
     return {dim: draw(indexer) for dim, indexer in concrete_indexers.items()}
 
 
+@st.composite
+def vectorized_indexers(draw, sizes):
+    max_size = max(sizes.values())
+    shape = draw(st.integers(min_value=1, max_value=max_size))
+    dtypes = npst.integer_dtypes()
+
+    indexers = {
+        dim: npst.arrays(
+            dtypes, shape, elements={"min_value": -size, "max_value": size - 1}
+        )
+        for dim, size in sizes.items()
+    }
+
+    return {
+        dim: xr.Variable("points", draw(indexer)) for dim, indexer in indexers.items()
+    }
+
+
 class IndexingTests(DuckArrayTestMixin):
     @property
     def orthogonal_indexer_types(self):
@@ -78,3 +97,17 @@ def test_variable_isel_orthogonal(self, data):
 
         assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}"
         self.assert_equal(actual, expected)
+
+    @given(st.data())
+    def test_variable_isel_vectorized(self, data):
+        variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
+        idx = data.draw(vectorized_indexers(variable.sizes))
+
+        with self.expected_errors("isel_vectorized", variable=variable):
+            actual = variable.isel(idx).data
+
+            raw_indexers = {dim: idx.get(dim, slice(None)) for dim in variable.dims}
+            expected = variable.data[*raw_indexers.values()]
+
+        assert isinstance(actual, self.array_type), f"wrong type: {type(actual)}"
+        self.assert_equal(actual, expected)