Skip to content

Commit 0172612

Browse files
d-v-bd-v-b
committed
fix(zarr-indexing): replace runtime assert statements with explicit checks
Asserts are stripped under `python -O`. The sixteen in zarr-indexing's runtime code narrowed types or guarded internal invariants rather than validating input, so most are restructured away: `inverse` validates and builds in one pass, the correlated-block mask starts all-true instead of `None`, the lowering paths branch on the map type instead of on membership in a dims list, the vindex path keeps its slices in a typed dict, and the affine writer returns early on an index array. The remainder become `RuntimeError`s for states the public API cannot reach, via a `_finite` helper in the message layer and an explicit check in `_indexed_set`. The `packages/zarr-indexing/src/**` exemption from ruff's S101 rule is removed so the package is held to the same standard as `src/zarr`. Assisted-by: ClaudeCode:claude-fable-5-1
1 parent 9f0ae4e commit 0172612

7 files changed

Lines changed: 97 additions & 82 deletions

File tree

‎packages/zarr-indexing/changes/4364.misc.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
Removed every `assert` statement from the package's runtime code and dropped
2+
the `zarr-indexing` exemption from the repo-wide ruff `S101` rule, so none can
3+
return. Asserts are stripped under `python -O`; the ones here narrowed types
4+
or guarded internal invariants rather than validating input, so most were
5+
restructured away (branching on the map type, carrying the narrowed value in a
6+
local) and the remainder became explicit `RuntimeError`s for states the
7+
public API cannot reach. No user-visible behavior changes.

‎packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -637,7 +637,14 @@ def _strided_set(
637637

638638
def _indexed_set(out_dim: int, m: ArrayMap, dg: DimensionGridLike) -> IndexedSet:
639639
dependent = m.dependent_axis
640-
assert dependent is not None
640+
if dependent is None:
641+
# `_partition_transform` sends every map that varies over no axis down
642+
# the joint (general) path; an orthogonal map always has its one axis.
643+
raise RuntimeError(
644+
f"output dimension {out_dim} has an index array of shape "
645+
f"{m.index_array.shape} that varies over no input axis; it cannot be "
646+
"partitioned as an orthogonal map"
647+
)
641648
flat = m.index_array.reshape(-1)
642649
n = int(flat.size)
643650
storage = checked_affine(m.offset, m.stride, flat)

‎packages/zarr-indexing/src/zarr_indexing/messages.py‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,18 @@ def _bound_is_implicit(bound: int | str | list[int | str]) -> bool:
203203
return isinstance(bound, list)
204204

205205

206+
def _finite(value: int | str, where: str) -> int:
207+
"""The integer behind a validated `index-value` that is not a sentinel.
208+
209+
Every value reaching the desugaring helpers has passed `_check_index_value`,
210+
so a string here is one of the two sentinels, which each helper handles
211+
before calling this. Anything else is an internal error, not bad input.
212+
"""
213+
if isinstance(value, int):
214+
return value
215+
raise RuntimeError(f"{where} is {value!r}; expected an integer or an infinity sentinel")
216+
217+
206218
def _ext_key(value: int | str) -> tuple[int, int]:
207219
"""A sort key giving the extended-integer order `-inf < n < +inf` exactly.
208220
@@ -213,8 +225,7 @@ def _ext_key(value: int | str) -> tuple[int, int]:
213225
return (0, 0)
214226
if value == "+inf":
215227
return (2, 0)
216-
assert isinstance(value, int)
217-
return (1, value)
228+
return (1, _finite(value, "bound"))
218229

219230

220231
def _rewrap(value: int | str, *, implicit: bool) -> int | str | list[int | str]:
@@ -276,10 +287,10 @@ def _resolve_upper_bound(
276287
The implicit/explicit bracket travels with the extent-bearing field (the
277288
upper bound, or `shape`), matching the spec's `[n]`-bracket convention.
278289
"""
279-
if upper_field is None:
290+
if upper_raw is None:
291+
# No upper-bound field was supplied (`upper_field` is None as well).
280292
return [["+inf"] for _ in range(rank)]
281293

282-
assert upper_raw is not None
283294
if kind_of == "exclusive":
284295
return list(upper_raw)
285296

@@ -316,8 +327,7 @@ def _checked_i64(value: int, where: str) -> int:
316327
def _inclusive_to_exclusive(value: int | str, where: str) -> int | str:
317328
if value == "+inf" or value == "-inf":
318329
return value
319-
assert isinstance(value, int)
320-
return _checked_i64(value + 1, f"{where} converted to an exclusive bound")
330+
return _checked_i64(_finite(value, where) + 1, f"{where} converted to an exclusive bound")
321331

322332

323333
def _shape_to_exclusive(min_value: int | str, shape_value: int | str, where: str) -> int | str:
@@ -330,9 +340,10 @@ def _shape_to_exclusive(min_value: int | str, shape_value: int | str, where: str
330340
return "+inf"
331341
if min_value == "-inf":
332342
return "-inf"
333-
assert isinstance(min_value, int)
334-
assert isinstance(shape_value, int)
335-
return _checked_i64(min_value + shape_value, f"{where} added to its inclusive_min")
343+
return _checked_i64(
344+
_finite(min_value, f"inclusive_min for {where}") + _finite(shape_value, where),
345+
f"{where} added to its inclusive_min",
346+
)
336347

337348

338349
def _validate_domain(inclusive_min: list[Any], exclusive_max: list[Any], *, prefix: str) -> None:

‎packages/zarr-indexing/src/zarr_indexing/reader.py‎

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -379,14 +379,14 @@ def _lower_orthogonal(array: Any, transform: IndexTransform) -> Any:
379379
if isinstance(m, ConstantMap):
380380
selection.append(m.offset)
381381
continue
382-
if out_dim in gathered:
383-
selection.append(slice(None))
384-
else:
385-
assert isinstance(m, DimensionMap)
382+
if isinstance(m, DimensionMap) and out_dim not in gathered:
386383
d = m.input_dimension
387384
lo = transform.domain.inclusive_min[d]
388385
hi = transform.domain.exclusive_max[d]
389386
selection.append(slice(m.offset + m.stride * lo, m.offset + m.stride * hi, m.stride))
387+
else:
388+
# Every ArrayMap, and every non-positive-stride DimensionMap, was gathered above.
389+
selection.append(slice(None))
390390
if isinstance(m, ArrayMap):
391391
axis = m.dependent_axis
392392
if axis is None:
@@ -414,15 +414,14 @@ def _lower_general(array: Any, transform: IndexTransform) -> Any:
414414
flat axis with row-major strides, and a single `take` collects them.
415415
"""
416416
outputs = transform.output
417-
correlated_dims = [d for d, m in enumerate(outputs) if isinstance(m, ArrayMap)]
417+
correlated = [(d, m) for d, m in enumerate(outputs) if isinstance(m, ArrayMap)]
418+
correlated_dims = [d for d, _ in correlated]
418419

419420
slice_input_dims = {m.input_dimension for m in outputs if isinstance(m, DimensionMap)}
420421
broadcast_axes = [d for d in range(transform.input_rank) if d not in slice_input_dims]
421422
broadcast_shape = tuple(transform.domain.shape[d] for d in broadcast_axes)
422423

423-
for d in correlated_dims:
424-
arr_map = outputs[d]
425-
assert isinstance(arr_map, ArrayMap)
424+
for _, arr_map in correlated:
426425
# The axes the array varies over (its non-singleton axes; see
427426
# transform._array_map_dependency_axes) must all live in the block.
428427
dependency = (axis for axis, size in enumerate(arr_map.index_array.shape) if size > 1)
@@ -454,16 +453,14 @@ def _lower_general(array: Any, transform: IndexTransform) -> Any:
454453
if isinstance(m, ConstantMap):
455454
selection.append(m.offset)
456455
continue
457-
if out_dim in correlated_dims:
456+
if isinstance(m, ArrayMap):
458457
selection.append(slice(None))
459458
correlated_positions.append(axis)
460459
elif out_dim in gathered:
461460
selection.append(slice(None))
462461
residual_positions.append(axis)
463-
assert isinstance(m, DimensionMap)
464462
residual_axis_dims.append(m.input_dimension)
465463
else:
466-
assert isinstance(m, DimensionMap)
467464
d = m.input_dimension
468465
lo = transform.domain.inclusive_min[d]
469466
hi = transform.domain.exclusive_max[d]
@@ -487,10 +484,9 @@ def _lower_general(array: Any, transform: IndexTransform) -> Any:
487484
flat_index = np.zeros(math.prod(broadcast_shape), dtype=np.intp)
488485
stride = 1
489486
for position in range(n_corr - 1, -1, -1):
490-
m = outputs[correlated_dims[position]]
491-
assert isinstance(m, ArrayMap)
487+
_, corr_map = correlated[position]
492488
flat_index = flat_index + (
493-
_correlated_map_coords(m, broadcast_axes, broadcast_shape, transform.input_rank)
489+
_correlated_map_coords(corr_map, broadcast_axes, broadcast_shape, transform.input_rank)
494490
* stride
495491
)
496492
stride *= corr_sizes[position]

‎packages/zarr-indexing/src/zarr_indexing/transform.py‎

Lines changed: 39 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929

3030
from __future__ import annotations
3131

32+
import math
3233
from dataclasses import dataclass
3334
from typing import TYPE_CHECKING, Any, Literal, cast
3435

@@ -398,43 +399,33 @@ def inverted(self) -> IndexTransform:
398399
f"{self.input_rank} and {self.output_rank}"
399400
)
400401

402+
# Nothing below touches `self`, so validating and building in one pass
403+
# leaves no partial state behind when a map is rejected.
401404
referenced: set[int] = set()
405+
inverse_min: list[int] = []
406+
inverse_max: list[int] = []
407+
inverse_output: dict[int, OutputIndexMap] = {}
402408
for output_dimension, output_map in enumerate(self.output):
403409
if isinstance(output_map, ArrayMap):
404410
raise ValueError( # noqa: TRY004 - valid map, invalid inverse
405411
f"cannot invert transform: output[{output_dimension}] is an ArrayMap"
406412
)
407-
if isinstance(output_map, DimensionMap):
408-
if output_map.stride not in (-1, 1):
409-
raise ValueError(
410-
"cannot invert transform: DimensionMap stride must be +1 or -1, "
411-
f"got {output_map.stride} for output[{output_dimension}]"
412-
)
413-
if output_map.input_dimension in referenced:
414-
raise ValueError(
415-
"cannot invert transform: input dimension "
416-
f"{output_map.input_dimension} is referenced more than once"
417-
)
418-
referenced.add(output_map.input_dimension)
419-
420-
for input_dimension, extent in enumerate(self.domain.shape):
421-
if input_dimension not in referenced and extent != 1:
422-
raise ValueError(
423-
"cannot invert transform: unreferenced input dimension "
424-
f"{input_dimension} has extent {extent}, not 1"
425-
)
426-
427-
inverse_min: list[int] = []
428-
inverse_max: list[int] = []
429-
inverse_output: dict[int, OutputIndexMap] = {}
430-
for output_dimension, output_map in enumerate(self.output):
431413
if isinstance(output_map, ConstantMap):
432414
inverse_min.append(output_map.offset)
433415
inverse_max.append(output_map.offset + 1)
434416
continue
435-
436-
assert isinstance(output_map, DimensionMap)
417+
if output_map.stride not in (-1, 1):
418+
raise ValueError(
419+
"cannot invert transform: DimensionMap stride must be +1 or -1, "
420+
f"got {output_map.stride} for output[{output_dimension}]"
421+
)
437422
input_dimension = output_map.input_dimension
423+
if input_dimension in referenced:
424+
raise ValueError(
425+
"cannot invert transform: input dimension "
426+
f"{input_dimension} is referenced more than once"
427+
)
428+
referenced.add(input_dimension)
438429
lower = self.domain.inclusive_min[input_dimension]
439430
upper = self.domain.exclusive_max[input_dimension]
440431
if output_map.stride == 1:
@@ -453,9 +444,17 @@ def inverted(self) -> IndexTransform:
453444
stride=-1,
454445
)
455446

456-
for input_dimension, lower in enumerate(self.domain.inclusive_min):
457-
if input_dimension not in referenced:
458-
inverse_output[input_dimension] = ConstantMap(lower)
447+
for input_dimension, (lower, extent) in enumerate(
448+
zip(self.domain.inclusive_min, self.domain.shape, strict=True)
449+
):
450+
if input_dimension in referenced:
451+
continue
452+
if extent != 1:
453+
raise ValueError(
454+
"cannot invert transform: unreferenced input dimension "
455+
f"{input_dimension} has extent {extent}, not 1"
456+
)
457+
inverse_output[input_dimension] = ConstantMap(lower)
459458

460459
return IndexTransform(
461460
domain=IndexDomain(tuple(inverse_min), tuple(inverse_max)),
@@ -1164,14 +1163,12 @@ def _intersect_general(
11641163
broadcast_shape = block.broadcast_shape
11651164

11661165
# Joint bounds mask over the broadcast block.
1167-
combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None
1166+
combined = np.ones(math.prod(broadcast_shape), dtype=np.bool_)
11681167
for out_dim in correlated_dims:
11691168
storage = block.flat_storage[out_dim]
11701169
lo = output_domain.inclusive_min[out_dim]
11711170
hi = output_domain.exclusive_max[out_dim]
1172-
mask = (storage >= lo) & (storage < hi)
1173-
combined = mask if combined is None else (combined & mask)
1174-
assert combined is not None
1171+
combined &= (storage >= lo) & (storage < hi)
11751172
surviving = np.flatnonzero(combined).astype(np.intp)
11761173
if surviving.size == 0:
11771174
return None
@@ -1222,19 +1219,18 @@ def _intersect_general(
12221219
corr_shape = points_shape + (1,) * n_slice
12231220
new_output: list[OutputIndexMap] = []
12241221
for out_dim, m in enumerate(transform.output):
1225-
if out_dim in correlated_dims:
1226-
corr = cast("ArrayMap", m)
1222+
if isinstance(m, ArrayMap):
1223+
# `correlated_dims` is exactly the ArrayMap output dimensions.
12271224
new_output.append(
12281225
ArrayMap(
12291226
index_array=corr_values[out_dim].reshape(corr_shape),
1230-
offset=corr.offset,
1231-
stride=corr.stride,
1227+
offset=m.offset,
1228+
stride=m.stride,
12321229
)
12331230
)
12341231
elif isinstance(m, ConstantMap):
12351232
new_output.append(m)
12361233
else:
1237-
assert isinstance(m, DimensionMap)
12381234
new_output.append(
12391235
DimensionMap(
12401236
input_dimension=new_input_dim_of[m.input_dimension],
@@ -1780,7 +1776,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform:
17801776

17811777
# Separate array dims and slice dims
17821778
array_dims: list[int] = []
1783-
slice_dims: list[int] = []
1779+
slices: dict[int, slice] = {}
17841780
arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = []
17851781

17861782
for i, sel in enumerate(processed):
@@ -1791,7 +1787,8 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform:
17911787
array_dims.append(i)
17921788
arrays.append(sel)
17931789
else:
1794-
slice_dims.append(i)
1790+
slices[i] = sel
1791+
slice_dims = list(slices)
17951792

17961793
# Determine the shared domain without expanding the index arrays: their
17971794
# singleton axes describe independent dependencies during grid partitioning.
@@ -1800,12 +1797,10 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform:
18001797
# Slice dimensions (preserved-domain literal semantics, like basic indexing)
18011798
slice_dim_params: dict[int, tuple[int, int, int]] = {}
18021799
slice_bounds: list[tuple[int, int]] = []
1803-
for old_dim in slice_dims:
1804-
sel = processed[old_dim]
1805-
assert isinstance(sel, slice)
1800+
for old_dim, sl in slices.items():
18061801
lo = transform.domain.inclusive_min[old_dim]
18071802
hi = transform.domain.exclusive_max[old_dim]
1808-
start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi)
1803+
start, step, origin, size = _resolve_slice_ts(sl, old_dim, lo, hi)
18091804
slice_bounds.append((origin, origin + size))
18101805
slice_dim_params[old_dim] = (start, step, origin)
18111806

‎packages/zarr-indexing/src/zarr_indexing/writer.py‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,7 @@ def _write_affine(source: Any, transform: IndexTransform, broadcast: Any) -> boo
123123
output.input_dimension for output in transform.output if isinstance(output, DimensionMap)
124124
]
125125
if (
126-
any(isinstance(output, ArrayMap) for output in transform.output)
127-
or len(set(axes)) != len(axes)
126+
len(set(axes)) != len(axes)
128127
or any(shape[axis] != 1 for axis in range(len(shape)) if axis not in axes)
129128
or any(
130129
isinstance(output, DimensionMap) and output.stride == 0 for output in transform.output
@@ -134,19 +133,21 @@ def _write_affine(source: Any, transform: IndexTransform, broadcast: Any) -> boo
134133
selection: list[int | slice] = []
135134
reverse_axes: list[int] = []
136135
for output in transform.output:
136+
if isinstance(output, ArrayMap):
137+
# An index array is a gather, not one basic assignment.
138+
return False
137139
if isinstance(output, ConstantMap):
138140
selection.append(output.offset)
141+
continue
142+
axis = output.input_dimension
143+
start = output.offset + output.stride * transform.domain.inclusive_min[axis]
144+
if output.stride < 0:
145+
reverse_axes.append(axes.index(axis))
146+
selection.append(
147+
slice(start + output.stride * (shape[axis] - 1), start + 1, -output.stride)
148+
)
139149
else:
140-
assert isinstance(output, DimensionMap)
141-
axis = output.input_dimension
142-
start = output.offset + output.stride * transform.domain.inclusive_min[axis]
143-
if output.stride < 0:
144-
reverse_axes.append(axes.index(axis))
145-
selection.append(
146-
slice(start + output.stride * (shape[axis] - 1), start + 1, -output.stride)
147-
)
148-
else:
149-
selection.append(slice(start, start + output.stride * shape[axis], output.stride))
150+
selection.append(slice(start, start + output.stride * shape[axis], output.stride))
150151
unused = tuple(axis for axis in range(len(shape)) if axis not in axes)
151152
squeezed = np.squeeze(broadcast, axis=unused)
152153
remaining = sorted(axes)

‎pyproject.toml‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -402,8 +402,6 @@ ignore = [
402402
"**/tests/**" = ["S101"]
403403
"**/examples/**" = ["S101"]
404404
"**/testing/**" = ["S101"]
405-
# zarr-indexing still has asserts in runtime code; tracked as a separate change.
406-
"packages/zarr-indexing/src/**" = ["S101"]
407405

408406
[tool.mypy]
409407
files = ["src", "tests"]

0 commit comments

Comments
 (0)