docs(indexing): inherit current-contract docstrings from audit

Assisted-by: Codex:GPT-6

Author: d-v-b

packages/zarr-indexing/changes/4345.doc.md

@@ -1 +1,3 @@
 Correct indexing, reader, serialization, cache, and integration descriptions to match supported behavior; qualify NumPy/TensorStore compatibility and performance claims.
+
+Describe current contracts in source and test docstrings instead of narrating prior implementations. Clarify that immutable index coordinates do not snapshot source values.

packages/zarr-indexing/src/zarr_indexing/boundary.py

@@ -210,11 +210,9 @@ def split_scalar_axes(
     scalar with the advanced indices for the purpose of placing the broadcast
     result, so the two disagree when a scalar and an index array are separated:
     `a[0, ..., [1, 2]]` has shape `(2, 3)` for a `(2, 3, 4)` array, where
-    `a[0][..., [1, 2]]` has shape `(3, 2)`. The earlier claim here that they
-    always agree rested on `a[0, [1, 2], :]`, where the indices are adjacent and
-    they happen to. Scalar-first is the documented dialect (see the `lazy_array`
-    module docstring) — the divergence is deliberate, and this note exists so
-    that the correct end is not "fixed" later.
+    `a[0][..., [1, 2]]` has shape `(3, 2)`. Scalar-first processing is the
+    wrapper's indexing dialect; it does not implement NumPy's full advanced-axis
+    placement rules.
 
     Parameters
     ----------

packages/zarr-indexing/src/zarr_indexing/domain.py

@@ -82,12 +82,10 @@ def __post_init__(self) -> None:
     def _unchecked(
         cls, inclusive_min: tuple[int, ...], exclusive_max: tuple[int, ...]
     ) -> IndexDomain:
-        """Build an unlabeled domain from bounds the caller has already established.
+        """Build an unlabeled domain from validated bounds.
 
-        Skips `__post_init__`. For internal producers that derive bounds from an
-        already-valid domain — chunk resolution builds two domains per chunk,
-        and re-validating them was a measurable share of a plan's cost.
-        """
+        Skip __post_init__; internal callers must ensure equal bound lengths
+        and inclusive_min <= exclusive_max in every dimension."""
         domain = object.__new__(cls)
         object.__setattr__(domain, "inclusive_min", inclusive_min)
         object.__setattr__(domain, "exclusive_max", exclusive_max)
@@ -210,13 +208,9 @@ def narrow(self, selection: Any) -> IndexDomain:
         Raises
         ------
         BoundsCheckError
-            If a bound lies outside this domain. A slice bound used to be
-            clamped instead, so `narrow(slice(-3, None))` on `[0, 10)` quietly
-            returned the whole axis — reading as the NumPy spelling of "the last
-            three" and answering with something else — and `narrow(slice(20,
-            30))` returned a domain its own parent did not contain. The rest of
-            the algebra states no clamping and no negative wrapping as an
-            invariant and enforces it; this is the one place that did not.
+            If an integer index or explicit slice bound lies outside this domain.
+            Bounds are checked as literal coordinates, without clamping or
+            negative-index wrapping.
         """
         normalized = _normalize_selection(selection, self.ndim)
         new_inclusive_min: list[int] = []

packages/zarr-indexing/src/zarr_indexing/lazy_array.py

@@ -898,9 +898,6 @@ def with_parts(self, parts: Sequence[int]) -> LazyArray:
         [`with_parts_per_axis`][zarr_indexing.lazy_array.LazyArray.with_parts_per_axis],
         and to read in one pass see
         [`unpartitioned`][zarr_indexing.lazy_array.LazyArray.unpartitioned].
-        The three were one parameter whose meaning was decided by inspecting the
-        type of what it was given, which left no way to ask for one of them and
-        be told when you had spelled it wrong.
 
         Parameters
         ----------

packages/zarr-indexing/src/zarr_indexing/output_map.py

@@ -194,15 +194,12 @@ class ArrayMap:
     """Multiplier applied to each `index_array` value before `offset` is added."""
 
     def __post_init__(self) -> None:
-        """Own the index array and expose it read-only.
-
-        A map is frozen, but the array inside it was not: reaching through a
-        view's transform to `index_array[0] = 9` silently changed what the view
-        returned, in a package whose whole contract is that a view is a
-        description of a read and resolving it twice answers alike. Owning the
-        array also prevents the caller from changing the contents behind the
-        read-only view, which would invalidate this value object's hash.
-        """
+        """Own an immutable snapshot of the integer index coordinates.
+
+        The snapshot is backed by immutable bytes, so callers cannot modify it
+        or re-enable its WRITEABLE flag. Changes to the supplied array do not
+        change the map's coordinates or hash. This freezes the coordinate
+        mapping, not the source values read through it."""
         # Immutable bytes are the ultimate owner so callers cannot re-enable
         # the WRITEABLE flag, as they can on a read-only array that owns its
         # allocation. `asarray` also accepts the NumPy scalars that reach here
@@ -222,29 +219,21 @@ def __reduce__(self) -> tuple[object, tuple[object, int, int]]:
         )
 
     def _with_affine(self, offset: int, stride: int) -> ArrayMap:
-        """This map's coordinates under a different affine adjustment.
+        """Return a map with a different affine adjustment.
 
-        The frozen index array is shared rather than copied: it is already
-        owned by immutable bytes and read-only, so the ownership invariant
-        `__post_init__` establishes holds for the new map too. Chunk
-        resolution translates every restricted map once per chunk, and
-        re-copying the array there dominated the cost of small selections.
-        """
+        Share the immutable index array while replacing the offset and stride.
+        This preserves coordinate ownership without copying the array."""
         new = object.__new__(ArrayMap)
         object.__setattr__(new, "index_array", self.index_array)
         object.__setattr__(new, "offset", offset)
         object.__setattr__(new, "stride", stride)
         return new
 
     def __eq__(self, other: object) -> bool:
-        """Value equality, comparing index arrays element-wise.
+        """Compare offset, stride, array shape, and index values.
 
-        The generated `__eq__` compares them with `==`, whose result for two
-        arrays is an array — so asking whether two maps are equal raised
-        `ValueError: the truth value of an array ... is ambiguous`. `frozen=True`
-        reads as a promise that a value can be compared and hashed, and this is
-        what makes good on it.
-        """
+        Return a scalar boolean for another ArrayMap and NotImplemented for
+        other types."""
         if not isinstance(other, ArrayMap):
             return NotImplemented
         return (
@@ -255,11 +244,10 @@ def __eq__(self, other: object) -> bool:
         )
 
     def __hash__(self) -> int:
-        """Hashed by the array's contents, so equal maps hash alike.
+        """Hash the offset, stride, array shape, and index bytes.
 
-        The generated `__hash__` hashed the ndarray itself, which is unhashable;
-        a map could therefore not go in a set, or key a cache.
-        """
+        The immutable coordinate snapshot keeps the hash stable, and equal
+        maps have equal hashes."""
         return hash(
             (
                 self.offset,

packages/zarr-indexing/src/zarr_indexing/testing/strategies.py

@@ -70,16 +70,11 @@ def _basic_entry(size: int) -> st.SearchStrategy[Any]:
 
 
 def _orthogonal_entry(size: int) -> st.SearchStrategy[Any]:
-    """One axis of an `oindex` selection.
+    """Generate one axis of an orthogonal selection.
 
-    The slices carry a step and are free to stop early. Drawing them as
-    `slice(start, size)` alone meant no strided or reversed slice ever reached
-    `oindex`, and no orthogonal selection ever stopped short of the axis end.
-
-    An empty coordinate list is drawn too. It selects nothing, which is legal
-    and is exactly the shape that lost its axis on the way through JSON — but
-    with `min_size=1` no fancy selection was ever empty.
-    """
+    Include scalar coordinates, coordinate lists, boolean masks, and slices
+    with positive or negative steps and varying endpoints. Empty coordinate
+    lists and all-False masks exercise selections with zero-length axes."""
     coordinate = st.integers(-size, size - 1)
     return st.one_of(
         coordinate,
@@ -119,13 +114,10 @@ def masks(draw: st.DrawFn, shape: tuple[int, ...]) -> np.ndarray[Any, np.dtype[n
 
 
 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.
+    """Generate an all-False mask with the given shape.
 
-    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.
-    """
+    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_))
 
 

packages/zarr-indexing/tests/test_composition.py

@@ -440,11 +440,7 @@ def test_composing_an_inner_array_with_a_broadcast_axis_wider_than_one_cell() ->
 
 
 def test_composing_a_one_dimensional_inner_array_under_a_higher_rank_outer() -> None:
-    """The shortcut gated on the output rank but sized by the input rank.
-
-    A rank-2 outer therefore built a rank-1 array for a rank-2 domain, which the
-    engine's own invariant then rejected.
-    """
+    """Composed index arrays retain the full input rank of their domain."""
     outer = IndexTransform(
         domain=IndexDomain.from_shape((2, 3)),
         output=(DimensionMap(input_dimension=0, offset=1, stride=1),),

packages/zarr-indexing/tests/test_json.py

@@ -356,12 +356,7 @@ def _index_array_body(index_array: Any, rank: int = 1, extent: int = 2) -> Index
     ids=["floats", "mixed", "bools", "strings", "string", "scalar", "nulls"],
 )
 def test_a_non_integer_index_array_is_rejected(index_array: Any, detail: str) -> None:
-    """An `index_array` addresses output coordinates, so it must be integral.
-
-    Lowering a float array silently truncated it (`[0.9, 1.9]` selected cells 0
-    and 1), a bool array coerced to 0/1, and a string array leaked a raw NumPy
-    `ValueError` from the middle of the conversion.
-    """
+    """Index arrays require integer coordinates and reject float, bool, and string dtypes."""
     with pytest.raises(NdselError) as excinfo:
         IndexTransform.from_json(_index_array_body(index_array))
     assert excinfo.value.reason == "invalid_json"
@@ -546,13 +541,7 @@ def test_an_ambiguous_empty_index_array_is_rejected() -> None:
     ids=["float", "string", "bool", "non-string-label", "out-of-range"],
 )
 def test_a_malformed_domain_document_is_rejected(document: Any, reason: str, detail: str) -> None:
-    """The domain loader validates what the message layer validates.
-
-    Reading the keys directly was a second, undefended way into the same
-    objects: a bare `int()` truncated `3.9` to 3, coerced `"3"` and `True`, and
-    let a non-string label into a `tuple[str, ...]` — each building a domain
-    that was not the document's, and re-dumping as a different document.
-    """
+    """The domain loader enforces coordinate types, integer bounds, and string labels."""
     with pytest.raises(NdselError) as excinfo:
         IndexDomain.from_json(document)
     assert excinfo.value.reason == reason

packages/zarr-indexing/tests/test_lazy_array.py

@@ -276,11 +276,7 @@ def test_array_map_dependent_axis_reports_no_axis() -> None:
 
 
 def test_scalar_on_a_fancy_axis_collapses_to_a_constant() -> None:
-    """The degenerate-collapse rule: an all-singleton ArrayMap becomes a ConstantMap.
-
-    Without it the map keeps an `input_dimension` naming an axis the integer
-    index just removed, which after renumbering aliases a different axis.
-    """
+    """Scalar selection collapses an all-singleton ArrayMap to a ConstantMap."""
     view = IndexTransform.from_shape((7, 5)).oindex[np.array([3, 1]), slice(None)][0]
     assert view.output[0] == ConstantMap(offset=3)
     assert isinstance(view.output[1], DimensionMap)
@@ -926,13 +922,10 @@ def test_a_boolean_mask_composes_onto_a_fancy_view() -> None:
 
 
 def test_an_ellipsis_only_vindex_step_preserves_a_correlated_gather() -> None:
-    """Regression: a slice-only vindex step misread correlated maps as orthogonal.
+    """Ellipsis-only vectorized selection preserves correlated coordinates.
 
-    `vindex[...]` (and `vindex[..., scalar]`, whose remainder after the scalar
-    is split off is ellipsis-only) used to stamp each correlated map with its
-    block axis as an orthogonal binding. Two "orthogonal" maps then shared one
-    input axis, and the partition walk rejected its own transform mid-read.
-    """
+    This also applies after scalar indices have been split off. The selected
+    values must agree across partitionings."""
     base = np.arange(16).reshape(4, 4)
 
     for parts in (None, (2, 2), (4, 4), (1, 3)):
@@ -2435,15 +2428,7 @@ def test_a_masked_source_keeps_its_mask_under_every_partitioning(parts: Any) ->
 
 @pytest.mark.parametrize("parts", [None, (2, 2), (3, 4)])
 def test_a_masked_source_keeps_its_mask_when_the_view_is_empty(parts: Any) -> None:
-    """An empty result is still a result, and its type must not depend on the parts.
-
-    An empty view is answered without reading the source at all, and that
-    shortcut reached for the array namespace's own `empty` — which knows nothing
-    about masks — so an unpartitioned empty view came back a plain array while
-    the same view partitioned came back masked. No cells either way, so nothing
-    about the values changed; the caller just got a different type depending on
-    how the read had been divided.
-    """
+    """Empty views of masked sources return masked arrays for every partitioning."""
     data = np.ma.masked_greater(np.arange(12).reshape(3, 4), 7)
     got = repartition(LazyArray(data), parts).lazy[:, 2:2].result()
     assert isinstance(got, np.ma.MaskedArray), parts
@@ -2453,11 +2438,7 @@ def test_a_masked_source_keeps_its_mask_when_the_view_is_empty(parts: Any) -> No
 def test_a_large_array_without_dask_refuses_to_claim_equality(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """Above the digest limit the fallback must miss a cache rather than lie.
-
-    Two arrays differing in one element used to token identically, because the
-    fallback described the shape and dtype and gave up on the contents.
-    """
+    """Without Dask, arrays above the digest limit receive distinct fallback tokens."""
     import sys
 
     monkeypatch.setitem(sys.modules, "dask.base", None)
@@ -2521,15 +2502,7 @@ def test_a_zero_chunk_on_a_nonempty_axis_is_still_rejected() -> None:
 def test_the_coverage_count_agrees_with_numpy_for_reversed_selections(
     selection: tuple[Any, ...],
 ) -> None:
-    """The safety net behind `result()`'s coverage assertion, checked on its own.
-
-    `_out_selection_cell_count` sizes a partition's `out_selection` without
-    materializing it, and `result()` trusts that count to decide whether the
-    walk covered the view. Nothing pinned it for a reversed slice, so dropping
-    its `start <= stop` guard — or wrapping the subtraction in `abs()` — left
-    the suite green. A net nobody tests only matters once something else breaks,
-    which is exactly when it needs to be right.
-    """
+    """Coverage counts match NumPy selection sizes for reversed slices."""
     data = reference()
     view = LazyArray(data).with_parts((2, 2, 2)).lazy[selection]
     out_shape = view.shape
@@ -2556,26 +2529,14 @@ def test_the_coverage_count_agrees_with_numpy_for_reversed_selections(
 def test_the_coverage_count_matches_numpy_for_intervals_the_fast_path_declines(
     selection: tuple[Any, ...], out_shape: tuple[int, ...], expected: int
 ) -> None:
-    """The guard on `result()`'s safety net, exercised where the walk cannot reach it.
-
-    A partition walk only ever produces concrete forward in-bounds intervals, so
-    the guard that keeps everything else off the subtraction fast path is not
-    reachable through `parts()` at all — which is why removing it left the whole
-    suite green. It is the net's own contract, so it is checked directly.
-    """
+    """Coverage counts match NumPy for intervals outside the subtraction fast path."""
     counted = _out_selection_cell_count(selection, out_shape)
     assert counted == np.empty(out_shape)[selection].size
     assert counted == expected
 
 
 def test_a_zero_dimensional_index_array_drops_its_axis_like_a_scalar() -> None:
-    """`a[np.array(2), :]` is `a[2, :]` in NumPy, and now here too.
-
-    Only Python and NumPy integers counted as scalars, so a 0-d array fell
-    through to the fancy path and was widened into a length-1 index array —
-    keeping an axis NumPy drops. That was a third answer, agreeing with neither
-    NumPy nor eager zarr, which rejects it.
-    """
+    """Zero-dimensional integer arrays drop axes in orthogonal and vectorized selections."""
     data = np.arange(20).reshape(4, 5)
     for mode, expected in (
         ("oindex", data[np.array(2), :]),
@@ -2591,12 +2552,7 @@ def test_a_zero_dimensional_index_array_drops_its_axis_like_a_scalar() -> None:
 
 
 def test_a_multidimensional_array_in_an_orthogonal_selection_is_refused() -> None:
-    """The rule belongs to the selection, so the message speaks its vocabulary.
-
-    Left to the engine, this surfaced as a rank complaint about an `index_array`
-    the caller never wrote — the transform layer's words for a mistake made two
-    layers above it.
-    """
+    """Reject multidimensional orthogonal index arrays with a selection-level error."""
     with pytest.raises(IndexError, match="must be 1-dimensional"):
         LazyArray(np.arange(20).reshape(4, 5)).lazy.oindex[[[0, 1], [2, 3]], slice(None)]
 
@@ -2611,13 +2567,7 @@ def test_with_parts_rejects_a_bare_integer() -> None:
 
 
 def test_fancy_composition_over_an_empty_axis() -> None:
-    """Regression: composing fancy steps over an empty axis stays unpinned.
-
-    The empty-domain branch of `compose` produces index arrays that are
-    singleton on every non-empty axis; pinning one to an axis it merely
-    broadcasts along made a later basic step index a size-1 axis positionally
-    and raise, deep inside a legal chain.
-    """
+    """Fancy composition over an empty axis preserves shape through later selections."""
     base = np.empty((3, 0, 6), dtype=np.int64)
     view = LazyArray(base).lazy.oindex[[2, 1], :, [5, 0, 3]]
     assert view.shape == (2, 0, 3)

packages/zarr-indexing/tests/test_lazy_array_stateful.py

@@ -9,13 +9,9 @@
 declared. The zarr case runs a smaller budget: it reads through a store, and it
 exercises the same code paths.
 
-This replaces a seeded `_random_chain` sweep in `test_lazy_array` that read
-chained selections through `parts()`. That sweep did reach the states it was
-meant to, but a rank-0 correlated view was absorbed by a reshape in `result()`
-and mirrored into the sweep rather than read as a failure; asserting the
-documented assembly literally makes that impossible to paper over.
-`test_lazy_array` keeps its `result()`-based sweep, which is the deterministic
-cross-flavor coverage this does not attempt.
+The state machine checks both result() and explicit partition assembly against
+NumPy, including the shape of each partition's values. test_lazy_array also
+provides deterministic selection-chain coverage across source flavors.
 """
 
 from __future__ import annotations