better approach

Author: dcherian

src/zarr/testing/strategies.py

@@ -562,3 +562,78 @@ def chunk_paths(draw: st.DrawFn, ndim: int, numblocks: tuple[int, ...], subset:
     )
     subset_slicer = slice(draw(st.integers(min_value=0, max_value=ndim))) if subset else slice(None)
     return "/".join(map(str, blockidx[subset_slicer]))
+
+
+@st.composite
+def complex_chunk_grids(draw: st.DrawFn) -> RectilinearChunkGrid:
+    ndim = draw(st.integers(min_value=1, max_value=3))
+    nchunks = draw(st.integers(min_value=10, max_value=100))
+    dim_chunks = st.lists(
+        st.integers(min_value=1, max_value=10), unique=True, min_size=nchunks, max_size=nchunks
+    )
+    if draw(st.booleans()):
+        event("using RectilinearChunkGrid")
+        chunk_shapes = draw(st.lists(dim_chunks, min_size=ndim, max_size=ndim))
+        return RectilinearChunkGrid(chunk_shapes=chunk_shapes)
+
+    else:
+        event("using RectilinearChunkGrid (run length encoded)")
+        repeats = st.lists(
+            st.integers(min_value=1, max_value=20), min_size=nchunks, max_size=nchunks
+        )
+        chunk_shapes_rle = [
+            [[c, r] for c, r in zip(draw(dim_chunks), draw(repeats), strict=True)]
+            for _ in range(ndim)
+        ]
+        return RectilinearChunkGrid(chunk_shapes=chunk_shapes_rle)
+
+
+@st.composite
+def complex_chunked_arrays(
+    draw: st.DrawFn,
+    *,
+    stores: st.SearchStrategy[StoreLike] = stores,
+) -> Array:
+    store = draw(stores, label="store")
+    chunks = draw(complex_chunk_grids(), label="chunk grid")
+    assert isinstance(chunks, RectilinearChunkGrid)
+    shape = tuple(x[-1] for x in chunks._cumulative_sizes)
+    nparray = draw(numpy_arrays(shapes=st.just(shape)), label="array data")
+    root = zarr.open_group(store, mode="w")
+
+    a = root.create_array(
+        "/foo",
+        shape=nparray.shape,
+        chunks=chunks,
+        shards=None,
+        dtype=nparray.dtype,
+        attributes={},
+        fill_value=None,
+        dimension_names=None,
+    )
+
+    assert isinstance(a, Array)
+    if a.metadata.zarr_format == 3:
+        assert a.fill_value is not None
+    assert nparray.shape == a.shape
+
+    # Verify chunks - for RegularChunkGrid check exact match
+    # For RectilinearChunkGrid, skip chunks check since it raises NotImplementedError
+    if isinstance(a.metadata.chunk_grid, RectilinearChunkGrid):
+        # Just verify the chunk_grid is set correctly
+        assert isinstance(a.metadata.chunk_grid, RectilinearChunkGrid)
+        # shards also raises NotImplementedError for RectilinearChunkGrid
+    else:
+        # For RegularChunkGrid, the chunks property returns the normalized chunk_shape
+        # which may differ from the input (e.g., (0,) becomes (1,) after normalization)
+        # We should compare against the actual chunk_grid's chunk_shape
+        from zarr.core.chunk_grids import RegularChunkGrid
+
+        assert isinstance(a.metadata.chunk_grid, RegularChunkGrid)
+        expected_chunks = a.metadata.chunk_grid.chunk_shape
+        assert expected_chunks == a.chunks
+
+    assert a.shards is None  # We don't use sharding with RectilinearChunkGrid
+
+    a[:] = nparray
+    return a

tests/test_properties.py

@@ -14,6 +14,7 @@
 import hypothesis.strategies as st
 from hypothesis import assume, given, settings
 
+from zarr import Array
 from zarr.abc.store import Store
 from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
@@ -22,6 +23,7 @@
     array_metadata,
     arrays,
     basic_indices,
+    complex_chunked_arrays,
     numpy_arrays,
     orthogonal_indices,
     simple_arrays,
@@ -106,11 +108,10 @@ def test_array_creates_implicit_groups(array):
 
 
 @pytest.mark.asyncio
-@settings(deadline=None)
+@settings(deadline=None, report_multiple_bugs=False)
 @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning")
-@given(data=st.data())
-async def test_basic_indexing(data: st.DataObject) -> None:
-    zarray = data.draw(simple_arrays())
+@given(data=st.data(), zarray=st.one_of([simple_arrays(), complex_chunked_arrays()]))
+async def test_basic_indexing(data: st.DataObject, zarray: Array) -> None:
     nparray = zarray[:]
     indexer = data.draw(basic_indices(shape=nparray.shape))
 
@@ -133,11 +134,18 @@ async def test_basic_indexing(data: st.DataObject) -> None:
 
 
 @pytest.mark.asyncio
-@given(data=st.data())
+@given(
+    data=st.data(),
+    zarray=st.one_of(
+        [
+            # integer_array_indices can't handle 0-size dimensions.
+            simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)),
+            complex_chunked_arrays(),
+        ]
+    ),
+)
 @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning")
-async def test_oindex(data: st.DataObject) -> None:
-    # integer_array_indices can't handle 0-size dimensions.
-    zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)))
+async def test_oindex(data: st.DataObject, zarray: Array) -> None:
     nparray = zarray[:]
     zindexer, npindexer = data.draw(orthogonal_indices(shape=nparray.shape))
 
@@ -165,11 +173,18 @@ async def test_oindex(data: st.DataObject) -> None:
 
 
 @pytest.mark.asyncio
-@given(data=st.data())
+@given(
+    data=st.data(),
+    zarray=st.one_of(
+        [
+            # integer_array_indices can't handle 0-size dimensions.
+            simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)),
+            complex_chunked_arrays(),
+        ]
+    ),
+)
 @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning")
-async def test_vindex(data: st.DataObject) -> None:
-    # integer_array_indices can't handle 0-size dimensions.
-    zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)))
+async def test_vindex(data: st.DataObject, zarray: Array) -> None:
     nparray = zarray[:]
     indexer = data.draw(
         npst.integer_array_indices(