zarr_indexing.testing.stateful

zarr_indexing.testing.stateful

A stateful property test for indexing an array through LazyArray.

ChainedIndexingStateMachine composes indexing steps onto a LazyArray wrapping your array — lazy[...], lazy.oindex[...], lazy.vindex[...], each step applied to the view the last one produced — while applying the same steps to a NumPy array holding the same values. After every step the view must still agree with that model three ways: its shape, its result(), and the assembly of its parts().

Point it at an array by subclassing and overriding make_source:

from zarr_indexing.testing import ChainedIndexingStateMachine, state_machine_test

class MyArrayIndexing(ChainedIndexingStateMachine):
    def make_source(self, data):
        array = my_format.create(shape=data.shape, dtype=data.dtype)
        array[:] = data
        return array

TestMyArrayIndexing = state_machine_test(MyArrayIndexing)

data, partitionings, and readers are class attributes; override any of them to widen or narrow what is drawn. The base class needs no make_source at all — left alone it wraps the NumPy array itself, which is a useful smoke test of this package but says nothing about yours.

What it is checking

The parts invariant is the one with teeth. Partition documents out[part.out_selection] = part.view.result() as the assembly procedure, so this checks that literally: a part's values must arrive at exactly the shape its out_selection addresses — not merely a shape that broadcasts into it — land there, and cover the view once. Checking through result() alone would prove only that result() is self-consistent.

The choose_reader rule draws a reader and applies it to the view, so the execution strategy becomes part of the chain. Every reader listed by a subclass must preserve the NumPy model for its source. The universal basic_reader is always exercised, even when a subclass lists only specialized readers; with no declared readers it is the sole strategy drawn.

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

DEFAULT_DATA module-attribute

DEFAULT_DATA = np.arange(7 * 5 * 4, dtype=np.int64).reshape(
    7, 5, 4
)

The values the source holds by default: distinct, so a misplaced cell shows.

DEFAULT_PARTITIONINGS module-attribute

DEFAULT_PARTITIONINGS: tuple[Any, ...] = (
    None,
    (2, 2, 2),
    (7, 5, 4),
    (3, 2, 3),
    ((3, 3, 1), (2, 2, 1), (3, 1)),
    (4, 3, 3),
)

Partitionings to read under: a single whole-array part, uniform boxes of several shapes (some of which do not divide the extent), and explicit per-axis sizes. Boxes that straddle whatever the source declares are deliberate — they cost extra I/O but must not change an answer.

DEFAULT_SETTINGS module-attribute

DEFAULT_SETTINGS = settings(
    max_examples=250,
    stateful_step_count=10,
    deadline=None,
    suppress_health_check=[
        HealthCheck.data_too_large,
        HealthCheck.filter_too_much,
        HealthCheck.too_slow,
    ],
)

Enough examples to find a defect reachable only through a narrow chain, at a few seconds per run when there is nothing to find. Every step is followed by checks that each materialize the whole view, so the budget buys examples rather than long chains — a chain runs out of axes to index within a few steps anyway.

filter_too_much is suppressed because a chain that reaches a rank-0 view leaves only repartition enabled, so a run that opens there is discarded.

__all__ module-attribute

__all__ = [
    "DEFAULT_DATA",
    "DEFAULT_PARTITIONINGS",
    "DEFAULT_SETTINGS",
    "ChainedIndexingStateMachine",
    "apply_selection",
    "outer_selection",
    "repartition",
    "state_machine_test",
]

ChainedIndexingStateMachine

Bases: RuleBasedStateMachine

Indexing steps composed onto one LazyArray, against NumPy as the model.

Subclass and override make_source to point it at your own array. See the module docstring for the shape of that subclass and for what the invariants check.

Attributes:

• data (Any) –

  The values the source holds, and the model every step is checked against. Any shape and dtype NumPy supports; every axis must be non-empty.

• partitionings (Sequence[Any]) –

  Drawn once per run, before any indexing: with_parts is a pure setter that carries through composition untouched and is read only when a view resolves, so choosing it up front reaches the same states choosing it mid-chain does, and spends the whole step budget on indexing.

• readers (Sequence[Reader] | None) –

  Execution strategies choose_reader may draw. Every listed reader must preserve the model for the source. basic_reader is always included; None means the reader already carried by the constructed view.

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

class ChainedIndexingStateMachine(RuleBasedStateMachine):
    """Indexing steps composed onto one `LazyArray`, against NumPy as the model.

    Subclass and override `make_source` to point it at your own array. See the
    module docstring for the shape of that subclass and for what the invariants
    check.

    Attributes
    ----------
    data
        The values the source holds, and the model every step is checked
        against. Any shape and dtype NumPy supports; every axis must be
        non-empty.
    partitionings
        Drawn once per run, before any indexing: `with_parts` is a pure setter
        that carries through composition untouched and is read only when a view
        resolves, so choosing it up front reaches the same states choosing it
        mid-chain does, and spends the whole step budget on indexing.
    readers
        Execution strategies `choose_reader` may draw. Every listed reader must
        preserve the model for the source. `basic_reader` is always included;
        `None` means the reader already carried by the constructed view.
    """

    data: ClassVar[Any] = DEFAULT_DATA
    partitionings: ClassVar[Sequence[Any]] = DEFAULT_PARTITIONINGS
    readers: ClassVar[Sequence[Reader] | None] = None

    def make_source(self, data: Any) -> Any:
        """Build the array under test, holding `data`.

        Called once per machine class and cached, not once per example: an
        example is cheap and a source may not be. The machine only ever reads,
        so the same object serves every run — but it must therefore not be
        mutated by anything else while the test runs.

        The default returns `data` itself, so an unsubclassed machine exercises
        this package against NumPy.
        """
        return data

    def __init__(self) -> None:
        super().__init__()
        cls = type(self)
        self.model: Any = np.asarray(cls.data)
        source = cls.__dict__.get(_SOURCE)
        if source is None:
            source = self.make_source(self.model)
            setattr(cls, _SOURCE, source)
        self.view = LazyArray(source)
        self.reader_choices = _reader_set(self.view, cls.readers)
        self.chain: list[tuple[str, Any]] = []

    def _indexable(self) -> bool:
        """Whether there is anything left to index.

        A rank-0 or empty view takes no further step — NumPy would reject one
        too — so the chain ends there, and the invariants keep checking.
        """
        return self.model.ndim > 0 and self.model.size > 0

    def _step(self, mode: SelectionMode, selection: tuple[Any, ...]) -> None:
        self.chain.append((mode, selection))
        self.model = apply_selection(self.model, selection, mode)
        if mode == "basic":
            self.view = self.view.lazy[selection]
        elif mode == "orthogonal":
            self.view = self.view.lazy.oindex[selection]
        else:
            self.view = self.view.lazy.vindex[selection]

    # -- rules --------------------------------------------------------------

    @initialize(data=st.data())
    def choose_partitioning(self, data: st.DataObject) -> None:
        """Fix how the read is broken up, before any indexing."""
        parts = data.draw(st.sampled_from(list(type(self).partitionings)))
        self.view = repartition(self.view, parts)
        self.chain.append(("parts", parts))

    @precondition(lambda self: self._indexable())
    @rule(data=st.data())
    def basic(self, data: st.DataObject) -> None:
        self._step("basic", data.draw(basic_selections(self.model.shape)))

    @precondition(lambda self: self._indexable())
    @rule(data=st.data())
    def orthogonal(self, data: st.DataObject) -> None:
        self._step("orthogonal", data.draw(orthogonal_selections(self.model.shape)))

    @precondition(lambda self: self._indexable())
    @rule(data=st.data())
    def vectorized(self, data: st.DataObject) -> None:
        self._step("vectorized", data.draw(vectorized_selections(self.model.shape)))

    @precondition(lambda self: self._indexable())
    @rule(data=st.data())
    def slices_only(self, data: st.DataObject) -> None:
        """An `oindex` step carrying only slices is not a fancy selection.

        It narrows the view's own axes and composes like basic indexing. Drawn
        as its own rule so that narrowing an existing index array by slices —
        a distinct code path from narrowing it with coordinates — stays
        exercised at full weight.
        """
        self._step("orthogonal", data.draw(slice_selections(self.model.shape)))

    @rule(data=st.data())
    def choose_reader(self, data: st.DataObject) -> None:
        """Read the rest of the chain through another conforming strategy."""
        reader = data.draw(st.sampled_from(list(self.reader_choices)))
        self.view = self.view.with_reader(reader)
        self.chain.append(("reader", type(reader).__qualname__))

    @precondition(lambda self: not self._indexable())
    @rule(data=st.data())
    def repartition(self, data: st.DataObject) -> None:
        """Re-box a chain that has run out of axes to index.

        Something must stay enabled once the view is rank-0 or empty, or
        Hypothesis has no move to make and abandons the run. Re-boxing is the
        useful thing to do there: it changes nothing the invariants may see, and
        a rank-0 view read through every partitioning is exactly the state a
        collapsed correlated selection reaches.
        """
        parts = data.draw(st.sampled_from(list(type(self).partitionings)))
        self.view = repartition(self.view, parts)
        self.chain.append(("parts", parts))

    # -- invariants ---------------------------------------------------------

    @invariant()
    def the_view_has_the_models_shape(self) -> None:
        assert self.view.shape == self.model.shape, self.chain

    @invariant()
    def result_matches_the_model(self) -> None:
        np.testing.assert_array_equal(
            np.asarray(self.view.result()), self.model, err_msg=str(self.chain)
        )

    @invariant()
    def parts_tile_the_view(self) -> None:
        """The documented assembly, run literally.

        Every part's values arrive at exactly the shape its `out_selection`
        addresses — not merely a shape that broadcasts into it — and together
        the parts cover the view once.
        """
        assembled = np.zeros(self.view.shape, dtype=self.view.dtype)
        hits = np.zeros(self.view.shape, dtype=np.int64)
        for part in self.view.parts():
            value = np.asarray(part.view.result())
            assert value.shape == assembled[part.out_selection].shape, (
                f"part {part.base_coords} carries {value.shape} for an out_selection "
                f"addressing {assembled[part.out_selection].shape}: {self.chain}"
            )
            # `is_complete` is what a consumer reads to decide it may take a
            # whole-box read and skip assembling anything, so a wrongly-`True`
            # one is the silent-corruption case. Asserted one way only: the flag
            # is documented as conservative, free to say `False` about a part it
            # does cover (a strided walk over a one-cell box, a fancy axis that
            # happens to enumerate everything), and only the claim to cover
            # everything has to be earned. Both quantities are already in hand
            # here, and nothing else in the suite compares them.
            if part.is_complete:
                box_cells = math.prod(stop - start for start, stop in part.box)
                assert value.size == box_cells, (
                    f"part {part.base_coords} reports is_complete but carries "
                    f"{value.size} of its box's {box_cells} cells: {self.chain}"
                )
            assembled[part.out_selection] = value
            np.add.at(hits, part.out_selection, 1)

        np.testing.assert_array_equal(assembled, self.model, err_msg=str(self.chain))
        np.testing.assert_array_equal(
            hits, np.ones(self.view.shape, dtype=np.int64), err_msg=str(self.chain)
        )

chain instance-attribute

chain: list[tuple[str, Any]] = []

data class-attribute

data: Any = DEFAULT_DATA

model instance-attribute

model: Any = np.asarray(cls.data)

partitionings class-attribute

partitionings: Sequence[Any] = DEFAULT_PARTITIONINGS

reader_choices instance-attribute

reader_choices = _reader_set(self.view, cls.readers)

readers class-attribute

readers: Sequence[Reader] | None = None

view instance-attribute

view = LazyArray(source)

__init__

__init__() -> None

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

def __init__(self) -> None:
    super().__init__()
    cls = type(self)
    self.model: Any = np.asarray(cls.data)
    source = cls.__dict__.get(_SOURCE)
    if source is None:
        source = self.make_source(self.model)
        setattr(cls, _SOURCE, source)
    self.view = LazyArray(source)
    self.reader_choices = _reader_set(self.view, cls.readers)
    self.chain: list[tuple[str, Any]] = []

basic

basic(data: DataObject) -> None

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

@precondition(lambda self: self._indexable())
@rule(data=st.data())
def basic(self, data: st.DataObject) -> None:
    self._step("basic", data.draw(basic_selections(self.model.shape)))

choose_partitioning

choose_partitioning(data: DataObject) -> None

Fix how the read is broken up, before any indexing.

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

@initialize(data=st.data())
def choose_partitioning(self, data: st.DataObject) -> None:
    """Fix how the read is broken up, before any indexing."""
    parts = data.draw(st.sampled_from(list(type(self).partitionings)))
    self.view = repartition(self.view, parts)
    self.chain.append(("parts", parts))

choose_reader

choose_reader(data: DataObject) -> None

Read the rest of the chain through another conforming strategy.

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

@rule(data=st.data())
def choose_reader(self, data: st.DataObject) -> None:
    """Read the rest of the chain through another conforming strategy."""
    reader = data.draw(st.sampled_from(list(self.reader_choices)))
    self.view = self.view.with_reader(reader)
    self.chain.append(("reader", type(reader).__qualname__))

make_source

make_source(data: Any) -> Any

Build the array under test, holding data.

Called once per machine class and cached, not once per example: an example is cheap and a source may not be. The machine only ever reads, so the same object serves every run — but it must therefore not be mutated by anything else while the test runs.

The default returns data itself, so an unsubclassed machine exercises this package against NumPy.

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

def make_source(self, data: Any) -> Any:
    """Build the array under test, holding `data`.

    Called once per machine class and cached, not once per example: an
    example is cheap and a source may not be. The machine only ever reads,
    so the same object serves every run — but it must therefore not be
    mutated by anything else while the test runs.

    The default returns `data` itself, so an unsubclassed machine exercises
    this package against NumPy.
    """
    return data

orthogonal

orthogonal(data: DataObject) -> None

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

@precondition(lambda self: self._indexable())
@rule(data=st.data())
def orthogonal(self, data: st.DataObject) -> None:
    self._step("orthogonal", data.draw(orthogonal_selections(self.model.shape)))

parts_tile_the_view

parts_tile_the_view() -> None

The documented assembly, run literally.

Every part's values arrive at exactly the shape its out_selection addresses — not merely a shape that broadcasts into it — and together the parts cover the view once.

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

@invariant()
def parts_tile_the_view(self) -> None:
    """The documented assembly, run literally.

    Every part's values arrive at exactly the shape its `out_selection`
    addresses — not merely a shape that broadcasts into it — and together
    the parts cover the view once.
    """
    assembled = np.zeros(self.view.shape, dtype=self.view.dtype)
    hits = np.zeros(self.view.shape, dtype=np.int64)
    for part in self.view.parts():
        value = np.asarray(part.view.result())
        assert value.shape == assembled[part.out_selection].shape, (
            f"part {part.base_coords} carries {value.shape} for an out_selection "
            f"addressing {assembled[part.out_selection].shape}: {self.chain}"
        )
        # `is_complete` is what a consumer reads to decide it may take a
        # whole-box read and skip assembling anything, so a wrongly-`True`
        # one is the silent-corruption case. Asserted one way only: the flag
        # is documented as conservative, free to say `False` about a part it
        # does cover (a strided walk over a one-cell box, a fancy axis that
        # happens to enumerate everything), and only the claim to cover
        # everything has to be earned. Both quantities are already in hand
        # here, and nothing else in the suite compares them.
        if part.is_complete:
            box_cells = math.prod(stop - start for start, stop in part.box)
            assert value.size == box_cells, (
                f"part {part.base_coords} reports is_complete but carries "
                f"{value.size} of its box's {box_cells} cells: {self.chain}"
            )
        assembled[part.out_selection] = value
        np.add.at(hits, part.out_selection, 1)

    np.testing.assert_array_equal(assembled, self.model, err_msg=str(self.chain))
    np.testing.assert_array_equal(
        hits, np.ones(self.view.shape, dtype=np.int64), err_msg=str(self.chain)
    )

repartition

repartition(data: DataObject) -> None

Re-box a chain that has run out of axes to index.

Something must stay enabled once the view is rank-0 or empty, or Hypothesis has no move to make and abandons the run. Re-boxing is the useful thing to do there: it changes nothing the invariants may see, and a rank-0 view read through every partitioning is exactly the state a collapsed correlated selection reaches.

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

@precondition(lambda self: not self._indexable())
@rule(data=st.data())
def repartition(self, data: st.DataObject) -> None:
    """Re-box a chain that has run out of axes to index.

    Something must stay enabled once the view is rank-0 or empty, or
    Hypothesis has no move to make and abandons the run. Re-boxing is the
    useful thing to do there: it changes nothing the invariants may see, and
    a rank-0 view read through every partitioning is exactly the state a
    collapsed correlated selection reaches.
    """
    parts = data.draw(st.sampled_from(list(type(self).partitionings)))
    self.view = repartition(self.view, parts)
    self.chain.append(("parts", parts))

result_matches_the_model

result_matches_the_model() -> None

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

@invariant()
def result_matches_the_model(self) -> None:
    np.testing.assert_array_equal(
        np.asarray(self.view.result()), self.model, err_msg=str(self.chain)
    )

slices_only

slices_only(data: DataObject) -> None

An oindex step carrying only slices is not a fancy selection.

It narrows the view's own axes and composes like basic indexing. Drawn as its own rule so that narrowing an existing index array by slices — a distinct code path from narrowing it with coordinates — stays exercised at full weight.

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

@precondition(lambda self: self._indexable())
@rule(data=st.data())
def slices_only(self, data: st.DataObject) -> None:
    """An `oindex` step carrying only slices is not a fancy selection.

    It narrows the view's own axes and composes like basic indexing. Drawn
    as its own rule so that narrowing an existing index array by slices —
    a distinct code path from narrowing it with coordinates — stays
    exercised at full weight.
    """
    self._step("orthogonal", data.draw(slice_selections(self.model.shape)))

the_view_has_the_models_shape

the_view_has_the_models_shape() -> None

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

@invariant()
def the_view_has_the_models_shape(self) -> None:
    assert self.view.shape == self.model.shape, self.chain

vectorized

vectorized(data: DataObject) -> None

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

@precondition(lambda self: self._indexable())
@rule(data=st.data())
def vectorized(self, data: st.DataObject) -> None:
    self._step("vectorized", data.draw(vectorized_selections(self.model.shape)))

apply_selection

apply_selection(
    array: Any,
    selection: tuple[Any, ...],
    mode: SelectionMode,
) -> Any

Apply a selection to a NumPy array in the given mode — the model a view is checked against.

NumPy's own semantics are basic and vectorized indexing, so only the orthogonal mode needs building (see outer_selection).

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

def apply_selection(array: Any, selection: tuple[Any, ...], mode: SelectionMode) -> Any:
    """Apply a selection to a NumPy array in the given mode — the model a view is checked against.

    NumPy's own semantics *are* basic and vectorized indexing, so only the
    orthogonal mode needs building (see `outer_selection`).
    """
    if mode == "orthogonal":
        return outer_selection(array, selection)
    return array[selection]

outer_selection

outer_selection(
    array: Any, selection: Sequence[Any]
) -> Any

Apply an orthogonal selection to a NumPy array: the outer product of its axes.

NumPy has no operator for this, so the model is built from numpy.ix_. Scalar integers are basic indices — NumPy applies them first and drops the axis — so they are peeled off before the outer product is formed.

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

def outer_selection(array: Any, selection: Sequence[Any]) -> Any:
    """Apply an orthogonal selection to a NumPy array: the outer product of its axes.

    NumPy has no operator for this, so the model is built from `numpy.ix_`.
    Scalar integers are basic indices — NumPy applies them first and drops the
    axis — so they are peeled off before the outer product is formed.
    """

    def is_scalar(sel: Any) -> bool:
        return isinstance(sel, (int, np.integer)) and not isinstance(sel, bool)

    scalars = tuple(sel if is_scalar(sel) else slice(None) for sel in selection)
    reduced = array[scalars]
    axes = [
        np.arange(size)[sel]
        for size, sel in zip(reduced.shape, [s for s in selection if not is_scalar(s)], strict=True)
    ]
    if len(axes) == 0:
        return reduced
    return reduced[np.ix_(*axes)]

repartition

repartition(view: LazyArray, parts: Any) -> LazyArray

Apply one of partitionings to a view.

The three partitioning spellings are three named methods, so a list holding a mix of them needs a dispatch somewhere. Choosing among them is what a test harness drawing from that list is doing, so it lives here rather than being pushed back into the public API as a type-inspecting parameter.

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

def repartition(view: LazyArray, parts: Any) -> LazyArray:
    """Apply one of `partitionings` to a view.

    The three partitioning spellings are three named methods, so a list holding
    a mix of them needs a dispatch somewhere. Choosing among them is what a test
    harness drawing from that list is doing, so it lives here rather than being
    pushed back into the public API as a type-inspecting parameter.
    """
    if parts is None:
        return view.unpartitioned()
    if any(isinstance(entry, Sequence) for entry in parts):
        return view.with_parts_per_axis(parts)
    return view.with_parts(parts)

state_machine_test

state_machine_test(
    machine: type[RuleBasedStateMachine],
    *,
    config: settings = DEFAULT_SETTINGS,
) -> Any

The pytest-collectable TestCase for a machine, with settings applied.

Hypothesis builds a fresh TestCase per state-machine class, so settings set on a base class do not reach a subclass's; this applies them where they land. Assign the result to a module-level name beginning with Test.

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

def state_machine_test(
    machine: type[RuleBasedStateMachine], *, config: settings = DEFAULT_SETTINGS
) -> Any:
    """The pytest-collectable `TestCase` for a machine, with settings applied.

    Hypothesis builds a fresh `TestCase` per state-machine class, so settings
    set on a base class do not reach a subclass's; this applies them where they
    land. Assign the result to a module-level name beginning with `Test`.
    """
    case = machine.TestCase
    case.settings = config
    return case