add oindex property test

Author: dcherian

src/zarr/testing/strategies.py

@@ -188,6 +188,34 @@ def basic_indices(draw: st.DrawFn, *, shape: tuple[int], **kwargs: Any) -> Any:
     )
 
 
+@st.composite  # type: ignore[misc]
+def orthogonal_indices(
+    draw: st.DrawFn, *, shape: tuple[int]
+) -> tuple[tuple[np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]:
+    """
+    Strategy that returns
+    (1) a tuple of integer arrays used for orthogonal indexing of Zarr arrays.
+    (2) an tuple of integer arrays that can be used for equivalent indexing of numpy arrays
+    """
+    zindexer = []
+    npindexer = []
+    ndim = len(shape)
+    for axis, size in enumerate(shape):
+        (idxr,) = draw(
+            npst.integer_array_indices(
+                shape=(size,), result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1)
+            )
+            # | npst.basic_indices(shape=(size,), allow_ellipsis=False)
+        )
+        zindexer.append(idxr)
+        if isinstance(idxr, np.ndarray):
+            newshape = [1] * ndim
+            newshape[axis] = idxr.size
+            idxr = idxr.reshape(newshape)
+            npindexer.append(idxr)
+    return tuple(zindexer), np.broadcast_arrays(*npindexer)
+
+
 def key_ranges(
     keys: SearchStrategy = node_names, max_size: int | None = None
 ) -> SearchStrategy[list[int]]:

tests/test_properties.py

@@ -6,9 +6,15 @@
 
 import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
-from hypothesis import given
+from hypothesis import given, settings
 
-from zarr.testing.strategies import arrays, basic_indices, numpy_arrays, zarr_formats
+from zarr.testing.strategies import (
+    arrays,
+    basic_indices,
+    numpy_arrays,
+    orthogonal_indices,
+    zarr_formats,
+)
 
 
 @given(data=st.data(), zarr_format=zarr_formats)
@@ -32,6 +38,18 @@ def test_basic_indexing(data: st.DataObject) -> None:
     assert_array_equal(nparray, zarray[:])
 
 
+@settings(report_multiple_bugs=False)
+@given(data=st.data())
+def test_oindex(data: st.DataObject) -> None:
+    # integer_array_indices can't handle 0-size dimensions.
+    zarray = data.draw(arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)))
+    nparray = zarray[:]
+
+    zindexer, npindexer = data.draw(orthogonal_indices(shape=nparray.shape))
+    actual = zarray.oindex[zindexer]
+    assert_array_equal(nparray[npindexer], actual)
+
+
 @given(data=st.data())
 def test_vindex(data: st.DataObject) -> None:
     # integer_array_indices can't handle 0-size dimensions.