More proptest updates

Author: dcherian

src/zarr/testing/strategies.py

@@ -233,22 +233,41 @@ def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[
     for size in shape:
         assert size > 0
         if size > 1:
-            # Limit max chunks to 20 to avoid performance issues with large chunk grids
-            max_chunks = min(size - 1, 20)
-            nchunks = draw(st.integers(min_value=1, max_value=max_chunks))
-            dividers = sorted(
-                draw(
-                    st.lists(
-                        st.integers(min_value=1, max_value=size - 1),
-                        min_size=nchunks - 1,
-                        max_size=nchunks - 1,
-                        unique=True,
+            mode = draw(st.sampled_from(["expanded", "rle"]))
+            if mode == "expanded":
+                event("rectilinear expanded")
+                # Limit max chunks to 20 to avoid performance issues with large chunk grids
+                max_chunks = min(size - 1, 20)
+                nchunks = draw(st.integers(min_value=1, max_value=max_chunks))
+                dividers = sorted(
+                    draw(
+                        st.lists(
+                            st.integers(min_value=1, max_value=size - 1),
+                            min_size=nchunks - 1,
+                            max_size=nchunks - 1,
+                            unique=True,
+                        )
                     )
                 )
-            )
-            chunk_shapes.append(
-                [a - b for a, b in zip(dividers + [size], [0] + dividers, strict=False)]
-            )
+                chunk_shapes.append(
+                    [a - b for a, b in zip(dividers + [size], [0] + dividers, strict=False)]
+                )
+            else:
+                # RLE mode: uniform chunks with optional remainder
+                max_chunk_size = min(size, 20)
+                chunk_size = draw(st.integers(min_value=1, max_value=max_chunk_size))
+                n_full = size // chunk_size
+                remainder = size % chunk_size
+                chunks_list = [chunk_size] * n_full
+                if remainder > 0:
+                    chunks_list.append(remainder)
+                # Optionally shuffle to create non-contiguous duplicate patterns
+                if draw(st.booleans()):
+                    event("rectilinear rle shuffled")
+                    chunks_list = draw(st.permutations(chunks_list))
+                else:
+                    event("rectilinear rle")
+                chunk_shapes.append(list(chunks_list))
         else:
             chunk_shapes.append([1])
     return chunk_shapes

tests/test_properties.py

@@ -24,6 +24,7 @@
     array_metadata,
     arrays,
     basic_indices,
+    chunk_grids,
     numpy_arrays,
     orthogonal_indices,
     simple_arrays,
@@ -373,3 +374,14 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N
         assert serialized_complex_float_is_valid(asdict_dict["fill_value"])
     elif dtype_native.kind in ("M", "m") and np.isnat(meta.fill_value):
         assert asdict_dict["fill_value"] == -9223372036854775808
+
+
+@given(
+    shape=npst.array_shapes(min_dims=1, max_dims=4, min_side=1, max_side=100),
+    data=st.data(),
+)
+@settings(max_examples=200)
+def test_chunk_grid_roundtrip(shape: tuple[int, ...], data: st.DataObject) -> None:
+    grid = data.draw(chunk_grids(shape=shape))
+    roundtripped = type(grid).from_dict(grid.to_dict())
+    assert roundtripped == grid