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 — an index tuple with one entry per axis, in the spelling its mode expects. 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])

Every axis of shape must be non-empty: a selection over an axis of extent 0 has no coordinates to draw. Filter or narrow the shape before calling.

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_]]]

The all-False mask over shape — a fancy selection that empties the view.

Split out from masks, which forces a cell True so a chain has something left to index at the next step. Drawn on its own because an empty fancy selection is a shape the code paths treat separately, and nothing generated one.

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

def empty_masks(shape: tuple[int, ...]) -> st.SearchStrategy[np.ndarray[Any, np.dtype[np.bool_]]]:
    """The all-False mask over `shape` — a fancy selection that empties the view.

    Split out from `masks`, which forces a cell True so a chain has something
    left to index at the next step. Drawn on its own because an empty fancy
    selection is a shape the code paths treat separately, and nothing generated
    one.
    """
    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 — so it is legal after a fancy step, where genuine coordinates are not. 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 — so it is legal after a fancy step, where
    genuine coordinates are not. 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)