zarr_indexing.testing.strategies

zarr_indexing.testing.strategies

Hypothesis strategies for the selections LazyArray accepts.

Each strategy takes the shape of the array being indexed and generates one selection for it, in the spelling its mode expects. Some vectorized selections use a partial coordinate tuple or a mask with an ellipsis. They are the generators behind ChainedIndexingStateMachine and are exported on their own for a project that has its own test harness and wants only the hard part.

from hypothesis import given, strategies as st
from zarr_indexing.testing.strategies import basic_selections

@given(selection=basic_selections((7, 5, 4)))
def test_my_array_slices_like_numpy(selection):
    assert_array_equal(my_array[selection], reference[selection])

Coordinate-drawing strategies require non-empty axes; vectorized selections also require positive rank. The empty_masks strategy can generate a mask for an empty shape because it does not draw element coordinates.

Requires the testing extra (pip install zarr-indexing[testing]).

__all__ module-attribute

__all__ = [
    "basic_selections",
    "empty_masks",
    "masks",
    "orthogonal_selections",
    "slice_selections",
    "vectorized_selections",
]

basic_selections

basic_selections(
    shape: tuple[int, ...],
) -> SearchStrategy[tuple[Any, ...]]

Basic selections: one scalar integer or slice per axis.

Slices run in both directions, including the two empty spellings — a forward slice whose stop precedes its start, and a backward one whose start is off the front of the axis.

Source code in src/zarr_indexing/testing/strategies.py

def basic_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]:
    """Basic selections: one scalar integer or slice per axis.

    Slices run in both directions, including the two empty spellings — a
    forward slice whose stop precedes its start, and a backward one whose start
    is off the front of the axis.
    """
    return _entries(shape, _basic_entry)

empty_masks

empty_masks(
    shape: tuple[int, ...],
) -> SearchStrategy[ndarray[Any, dtype[bool_]]]

Generate an all-False mask with the given shape.

This selects no elements. Unlike masks(), which includes a True cell, this strategy exercises empty fancy selections.

Source code in src/zarr_indexing/testing/strategies.py

def empty_masks(shape: tuple[int, ...]) -> st.SearchStrategy[np.ndarray[Any, np.dtype[np.bool_]]]:
    """Generate an all-False mask with the given shape.

    This selects no elements. Unlike masks(), which includes a True cell,
    this strategy exercises empty fancy selections."""
    return st.just(np.zeros(shape, dtype=np.bool_))

masks

masks(
    draw: DrawFn, shape: tuple[int, ...]
) -> ndarray[Any, dtype[bool_]]

Boolean masks over shape, each selecting at least one cell.

An all-False mask is legal but is a separate concern — it empties the view, and a chain of selections is more interesting when every step leaves something to index — so one cell is always forced True.

Source code in src/zarr_indexing/testing/strategies.py

@st.composite
def masks(draw: st.DrawFn, shape: tuple[int, ...]) -> np.ndarray[Any, np.dtype[np.bool_]]:
    """Boolean masks over `shape`, each selecting at least one cell.

    An all-False mask is legal but is a separate concern — it empties the view,
    and a chain of selections is more interesting when every step leaves
    something to index — so one cell is always forced True.
    """
    size = int(np.prod(shape))
    flags = np.array(draw(st.lists(st.booleans(), min_size=size, max_size=size)))
    flags[draw(st.integers(0, size - 1))] = True
    return flags.reshape(shape)

orthogonal_selections

orthogonal_selections(
    shape: tuple[int, ...],
) -> SearchStrategy[tuple[Any, ...]]

Orthogonal (oindex) selections: an outer product of per-axis choices.

Each axis draws a scalar, a coordinate list (unsorted, with duplicates), a boolean mask, or a slice.

Source code in src/zarr_indexing/testing/strategies.py

def orthogonal_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]:
    """Orthogonal (`oindex`) selections: an outer product of per-axis choices.

    Each axis draws a scalar, a coordinate list (unsorted, with duplicates), a
    boolean mask, or a slice.
    """
    return _entries(shape, _orthogonal_entry)

slice_selections

slice_selections(
    shape: tuple[int, ...],
) -> SearchStrategy[tuple[Any, ...]]

Selections of slices alone, for the oindex spelling that carries no coordinates.

Such a step is not a fancy selection — it narrows the view's own axes and composes like basic indexing. The starts reach past the origin, which is what distinguishes a step that walks an existing index array's dependency axes from one that walks its broadcast singletons.

Source code in src/zarr_indexing/testing/strategies.py

def slice_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]:
    """Selections of slices alone, for the `oindex` spelling that carries no coordinates.

    Such a step is not a fancy selection — it narrows the view's own axes and
    composes like basic indexing. The starts reach past the origin, which is what
    distinguishes a step that walks an existing index array's dependency axes
    from one that walks its broadcast singletons.
    """
    return _entries(shape, _slice_entry)

vectorized_selections

vectorized_selections(
    draw: DrawFn, shape: tuple[int, ...]
) -> tuple[Any, ...]

Vectorized (vindex) selections over a leading or trailing block of axes.

vindex is coordinate-only — it rejects a slice outright — so a partial selection names its axes by position: a leading block, or a trailing one reached through an ellipsis. Either a single boolean mask spanning the whole covered block, or one entry per axis, each a coordinate array or a scalar (a scalar being a basic index NumPy applies before the coordinates).

Source code in src/zarr_indexing/testing/strategies.py

@st.composite
def vectorized_selections(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[Any, ...]:
    """Vectorized (`vindex`) selections over a leading or trailing block of axes.

    `vindex` is coordinate-only — it rejects a slice outright — so a partial
    selection names its axes by position: a leading block, or a trailing one
    reached through an ellipsis. Either a single boolean mask spanning the whole
    covered block, or one entry per axis, each a coordinate array or a scalar
    (a scalar being a basic index NumPy applies before the coordinates).
    """
    ndim = len(shape)
    trailing = draw(st.booleans())
    count = draw(st.integers(1, ndim))
    axes = range(ndim - count, ndim) if trailing else range(count)
    sizes = [shape[axis] for axis in axes]

    entries: list[Any]
    if draw(st.booleans()):
        entries = [draw(masks(tuple(sizes)))]
    else:
        # The coordinate arrays share one shape, which is what makes the
        # selection correlated. That shape is not always one-dimensional: a
        # vectorized read of a (2, 3) block of points is an ordinary thing to
        # ask for and produces a result of that rank. Drawing only 1-D arrays
        # meant no rank-raising vindex was ever generated — and a length of 0
        # covers the empty case the same way `_orthogonal_entry` does.
        coordinate_shape = draw(
            st.one_of(
                st.integers(0, 4).map(lambda length: (length,)),
                st.tuples(st.integers(1, 2), st.integers(1, 3)),
            )
        )
        entries = [
            draw(
                st.one_of(
                    st.integers(-size, size - 1),
                    st.lists(
                        st.integers(-size, size - 1),
                        min_size=int(np.prod(coordinate_shape)),
                        max_size=int(np.prod(coordinate_shape)),
                    ).map(lambda values: np.array(values, dtype=np.intp).reshape(coordinate_shape)),
                )
            )
            for size in sizes
        ]
    return (Ellipsis, *entries) if trailing else tuple(entries)