Greatly improve property tests

Author: dcherian

src/zarr/testing/stateful.py

@@ -325,7 +325,7 @@ def __init__(self, store: Store) -> None:
     def init_store(self) -> None:
         self.store.clear()
 
-    @rule(key=zarr_keys, data=st.binary(min_size=0, max_size=MAX_BINARY_SIZE))
+    @rule(key=zarr_keys(), data=st.binary(min_size=0, max_size=MAX_BINARY_SIZE))
     def set(self, key: str, data: DataObject) -> None:
         note(f"(set) Setting {key!r} with {data}")
         assert not self.store.read_only
@@ -334,7 +334,7 @@ def set(self, key: str, data: DataObject) -> None:
         self.model[key] = data_buf
 
     @precondition(lambda self: len(self.model.keys()) > 0)
-    @rule(key=zarr_keys, data=st.data())
+    @rule(key=zarr_keys(), data=st.data())
     def get(self, key: str, data: DataObject) -> None:
         key = data.draw(
             st.sampled_from(sorted(self.model.keys()))
@@ -344,7 +344,7 @@ def get(self, key: str, data: DataObject) -> None:
         # to bytes here necessary because data_buf set to model in set()
         assert self.model[key] == store_value
 
-    @rule(key=zarr_keys, data=st.data())
+    @rule(key=zarr_keys(), data=st.data())
     def get_invalid_zarr_keys(self, key: str, data: DataObject) -> None:
         note("(get_invalid)")
         assume(key not in self.model)
@@ -408,7 +408,7 @@ def is_empty(self) -> None:
         # make sure they either both are or both aren't empty (same state)
         assert self.store.is_empty("") == (not self.model)
 
-    @rule(key=zarr_keys)
+    @rule(key=zarr_keys())
     def exists(self, key: str) -> None:
         note("(exists)")
 

src/zarr/testing/strategies.py

@@ -1,10 +1,11 @@
+import math
 import sys
 from typing import Any, Literal
 
 import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
 import numpy as np
-from hypothesis import given, settings  # noqa: F401
+from hypothesis import event, given, settings  # noqa: F401
 from hypothesis.strategies import SearchStrategy
 
 import zarr
@@ -28,6 +29,16 @@
 )
 
 
+@st.composite  # type: ignore[misc]
+def keys(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> Any:
+    return draw(st.lists(node_names, min_size=1, max_size=max_num_nodes).map("/".join))
+
+
+@st.composite  # type: ignore[misc]
+def paths(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> Any:
+    return draw(st.just("/") | keys(max_num_nodes=max_num_nodes))
+
+
 def v3_dtypes() -> st.SearchStrategy[np.dtype]:
     return (
         npst.boolean_dtypes()
@@ -87,17 +98,19 @@ def clear_store(x: Store) -> Store:
 node_names = st.text(zarr_key_chars, min_size=1).filter(
     lambda t: t not in (".", "..") and not t.startswith("__")
 )
+short_node_names = st.text(zarr_key_chars, max_size=3, min_size=1).filter(
+    lambda t: t not in (".", "..") and not t.startswith("__")
+)
 array_names = node_names
 attrs = st.none() | st.dictionaries(_attr_keys, _attr_values)
-keys = st.lists(node_names, min_size=1).map("/".join)
-paths = st.just("/") | keys
 # st.builds will only call a new store constructor for different keyword arguments
 # i.e. stores.examples() will always return the same object per Store class.
 # So we map a clear to reset the store.
 stores = st.builds(MemoryStore, st.just({})).map(clear_store)
 compressors = st.sampled_from([None, "default"])
 zarr_formats: st.SearchStrategy[ZarrFormat] = st.sampled_from([3, 2])
-array_shapes = npst.array_shapes(max_dims=4, min_side=0)
+# We de-prioritize arrays having dim sizes 0, 1, 2
+array_shapes = npst.array_shapes(max_dims=4, min_side=3) | npst.array_shapes(max_dims=4, min_side=0)
 
 
 @st.composite  # type: ignore[misc]
@@ -174,11 +187,19 @@ def chunk_shapes(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[int, ...]:
         st.tuples(*[st.integers(min_value=0 if size == 0 else 1, max_value=size) for size in shape])
     )
     # 2. and now generate the chunks tuple
-    return tuple(
+    chunks = tuple(
         size // nchunks if nchunks > 0 else 0
         for size, nchunks in zip(shape, numchunks, strict=True)
     )
 
+    for c in chunks:
+        event("chunk size", c)
+
+    if any((c != 0 and s % c != 0) for s, c in zip(shape, chunks, strict=True)):
+        event("smaller last chunk")
+
+    return chunks
+
 
 @st.composite  # type: ignore[misc]
 def shard_shapes(
@@ -211,7 +232,7 @@ def arrays(
     shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes,
     compressors: st.SearchStrategy = compressors,
     stores: st.SearchStrategy[StoreLike] = stores,
-    paths: st.SearchStrategy[str | None] = paths,
+    paths: st.SearchStrategy[str | None] = paths(),  # noqa: B008
     array_names: st.SearchStrategy = array_names,
     arrays: st.SearchStrategy | None = None,
     attrs: st.SearchStrategy = attrs,
@@ -267,23 +288,56 @@ def arrays(
     return a
 
 
+@st.composite  # type: ignore[misc]
+def simple_arrays(
+    draw: st.DrawFn,
+    *,
+    shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes,
+) -> Any:
+    return draw(
+        arrays(
+            shapes=shapes,
+            paths=paths(max_num_nodes=2),
+            array_names=short_node_names,
+            attrs=st.none(),
+            compressors=st.sampled_from([None, "default"]),
+        )
+    )
+
+
 def is_negative_slice(idx: Any) -> bool:
     return isinstance(idx, slice) and idx.step is not None and idx.step < 0
 
 
+@st.composite  # type: ignore[misc]
+def end_slices(draw: st.DrawFn, *, shape: tuple[int]) -> Any:
+    """
+    A strategy that slices ranges that include the last chunk.
+    This is intended to stress-test handling of a possibly smaller last chunk.
+    """
+    slicers = []
+    for size in shape:
+        start = draw(st.integers(min_value=size // 2, max_value=size - 1))
+        length = draw(st.integers(min_value=0, max_value=size - start))
+        slicers.append(slice(start, start + length))
+    event("drawing end slice")
+    return tuple(slicers)
+
+
 @st.composite  # type: ignore[misc]
 def basic_indices(draw: st.DrawFn, *, shape: tuple[int], **kwargs: Any) -> Any:
     """Basic indices without unsupported negative slices."""
-    return draw(
-        npst.basic_indices(shape=shape, **kwargs).filter(
-            lambda idxr: (
-                not (
-                    is_negative_slice(idxr)
-                    or (isinstance(idxr, tuple) and any(is_negative_slice(idx) for idx in idxr))
-                )
+    strategy = npst.basic_indices(shape=shape, **kwargs).filter(
+        lambda idxr: (
+            not (
+                is_negative_slice(idxr)
+                or (isinstance(idxr, tuple) and any(is_negative_slice(idx) for idx in idxr))
             )
         )
     )
+    if math.prod(shape) >= 3:
+        strategy = end_slices(shape=shape) | strategy
+    return draw(strategy)
 
 
 @st.composite  # type: ignore[misc]

tests/test_indexing.py

@@ -425,6 +425,17 @@ def test_orthogonal_indexing_fallback_on_getitem_2d(
     np.testing.assert_array_equal(z[index], expected_result)
 
 
+def test_setitem_repeated_index():
+    array = zarr.array(data=np.zeros((4,)), chunks=(1,))
+    indexer = np.array([-1, -1, 0, 0])
+    array.oindex[(indexer,)] = [0, 1, 2, 3]
+    np.testing.assert_array_equal(array[:], np.array([3, 0, 0, 1]))
+
+    indexer = np.array([-1, 0, 0, -1])
+    array.oindex[(indexer,)] = [0, 1, 2, 3]
+    np.testing.assert_array_equal(array[:], np.array([2, 0, 0, 3]))
+
+
 Index = list[int] | tuple[slice | int | list[int], ...]
 
 
@@ -816,6 +827,25 @@ def test_set_orthogonal_selection_1d(store: StorePath) -> None:
         _test_set_orthogonal_selection(v, a, z, selection)
 
 
+def test_set_item_1d_last_two_chunks():
+    # regression test for GH2849
+    g = zarr.open_group("foo.zarr", zarr_format=3, mode="w")
+    a = g.create_array("bar", shape=(10,), chunks=(3,), dtype=int)
+    data = np.array([7, 8, 9])
+    a[slice(7, 10)] = data
+    np.testing.assert_array_equal(a[slice(7, 10)], data)
+
+    z = zarr.open_group("foo.zarr", mode="w")
+    z.create_array("zoo", dtype=float, shape=())
+    z["zoo"][...] = np.array(1)  # why doesn't [:] work?
+    np.testing.assert_equal(z["zoo"][()], np.array(1))
+
+    z = zarr.open_group("foo.zarr", mode="w")
+    z.create_array("zoo", dtype=float, shape=())
+    z["zoo"][...] = 1  # why doesn't [:] work?
+    np.testing.assert_equal(z["zoo"][()], np.array(1))
+
+
 def _test_set_orthogonal_selection_2d(
     v: npt.NDArray[np.int_],
     a: npt.NDArray[np.int_],

tests/test_properties.py

@@ -7,7 +7,7 @@
 
 import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
-from hypothesis import assume, given
+from hypothesis import assume, given, note
 
 from zarr.abc.store import Store
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
@@ -18,6 +18,7 @@
     basic_indices,
     numpy_arrays,
     orthogonal_indices,
+    simple_arrays,
     stores,
     zarr_formats,
 )
@@ -50,7 +51,7 @@ def test_array_creates_implicit_groups(array):
 
 @given(data=st.data())
 def test_basic_indexing(data: st.DataObject) -> None:
-    zarray = data.draw(arrays())
+    zarray = data.draw(simple_arrays())
     nparray = zarray[:]
     indexer = data.draw(basic_indices(shape=nparray.shape))
     actual = zarray[indexer]
@@ -65,7 +66,7 @@ def test_basic_indexing(data: st.DataObject) -> None:
 @given(data=st.data())
 def test_oindex(data: st.DataObject) -> None:
     # integer_array_indices can't handle 0-size dimensions.
-    zarray = data.draw(arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)))
+    zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)))
     nparray = zarray[:]
 
     zindexer, npindexer = data.draw(orthogonal_indices(shape=nparray.shape))
@@ -76,13 +77,14 @@ def test_oindex(data: st.DataObject) -> None:
     new_data = data.draw(npst.arrays(shape=st.just(actual.shape), dtype=nparray.dtype))
     nparray[npindexer] = new_data
     zarray.oindex[zindexer] = new_data
+    note((new_data, npindexer, nparray, zindexer, zarray[:]))
     assert_array_equal(nparray, zarray[:])
 
 
 @given(data=st.data())
 def test_vindex(data: st.DataObject) -> None:
     # integer_array_indices can't handle 0-size dimensions.
-    zarray = data.draw(arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)))
+    zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)))
     nparray = zarray[:]
 
     indexer = data.draw(