replace the orthogonal indexing test with the new strategy

Author: keewis

xarray_array_testing/indexing.py

@@ -1,73 +1,36 @@
 from contextlib import nullcontext
 
-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
 
 from xarray_array_testing.base import DuckArrayTestMixin
+from xarray_array_testing.strategies import orthogonal_indexers, vectorized_indexers
 
 
-def scalar_indexer(size):
-    return st.integers(min_value=-size, max_value=size - 1)
+def broadcast_orthogonal_indexers(indexers, sizes, *, xp):
+    def _broadcasting_shape(index, total):
+        return tuple(1 if i != index else -1 for i in range(total))
 
+    def _as_array(indexer, size):
+        if isinstance(indexer, slice):
+            return xp.asarray(range(*indexer.indices(size)), dtype="int64")
+        elif isinstance(indexer, int):
+            return xp.asarray(indexer, dtype="int64")
+        else:
+            return indexer
 
-def integer_array_indexer(size):
-    dtypes = npst.integer_dtypes()
-
-    return npst.arrays(
-        dtypes, size, elements={"min_value": -size, "max_value": size - 1}
-    )
-
-
-def indexers(size, indexer_types):
-    indexer_strategy_fns = {
-        "scalars": scalar_indexer,
-        "slices": st.slices,
-        "integer_arrays": integer_array_indexer,
+    indexer_arrays = {
+        dim: _as_array(indexer, sizes[dim]) for dim, indexer in indexers.items()
     }
-
-    bad_types = set(indexer_types) - indexer_strategy_fns.keys()
-    if bad_types:
-        raise ValueError(f"unknown indexer strategies: {sorted(bad_types)}")
-
-    # use the order of definition to prefer simpler strategies over more complex
-    # ones
-    indexer_strategies = [
-        strategy_fn(size)
-        for name, strategy_fn in indexer_strategy_fns.items()
-        if name in indexer_types
-    ]
-    return st.one_of(*indexer_strategies)
-
-
-@st.composite
-def orthogonal_indexers(draw, sizes, indexer_types):
-    # TODO: make use of `flatmap` and `builds` instead of `composite`
-    possible_indexers = {
-        dim: indexers(size, indexer_types) for dim, size in sizes.items()
-    }
-    concrete_indexers = draw(xrst.unique_subset_of(possible_indexers))
-    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}
+    broadcasted = xp.broadcast_arrays(
+        *(
+            xp.reshape(indexer, _broadcasting_shape(index, total=len(indexers)))
+            for index, indexer in enumerate(indexer_arrays.values())
         )
-        for dim, size in sizes.items()
-    }
+    )
 
-    return {
-        dim: xr.Variable("points", draw(indexer)) for dim, indexer in indexers.items()
-    }
+    return dict(zip(indexer_arrays.keys(), broadcasted))
 
 
 class IndexingTests(DuckArrayTestMixin):
@@ -81,19 +44,24 @@ def expected_errors(op, **parameters):
 
     @given(st.data())
     def test_variable_isel_orthogonal(self, data):
-        indexer_types = data.draw(
-            st.lists(self.orthogonal_indexer_types, min_size=1, unique=True)
-        )
         variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
-        idx = data.draw(orthogonal_indexers(variable.sizes, indexer_types))
+        idx = data.draw(
+            orthogonal_indexers(sizes=variable.sizes, min_dims=len(variable.dims))
+        )
 
-        with self.expected_errors(
-            "isel_orthogonal", variable=variable, indexer_types=indexer_types
-        ):
+        with self.expected_errors("isel_orthogonal", variable=variable, indexers=idx):
             actual = variable.isel(idx).data
 
-            raw_indexers = {dim: idx.get(dim, slice(None)) for dim in variable.dims}
-            expected = variable.data[*raw_indexers.values()]
+            sorted_dims = sorted(idx.keys(), key=variable.dims.index, reverse=True)
+            expected = variable.data
+            for dim in sorted_dims:
+                indexer = idx[dim]
+                axis = variable.get_axis_num(dim)
+                if isinstance(indexer, slice):
+                    indexer = self.xp.asarray(
+                        range(*indexer.indices(variable.sizes[dim])), dtype="int64"
+                    )
+                expected = self.xp.take(expected, indexer, axis=axis)
 
         assert isinstance(
             actual, self.array_type("orthogonal_indexing")

xarray_array_testing/strategies.py

@@ -1,11 +1,11 @@
 from collections.abc import Hashable
 from itertools import compress
 
-import hypothesis.extras.numpy as npst
+import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
 import numpy as np
 import xarray as xr
-from xr.testing.strategies import unique_subset_of
+from xarray.testing.strategies import unique_subset_of
 
 
 def _basic_indexers(size):