Fix typing errors in testing.stateful

Author: dstansby

.pre-commit-config.yaml

@@ -37,6 +37,7 @@ repos:
           - obstore>=0.5.1
           # Tests
           - pytest
+          - hypothesis
   - repo: https://github.com/scientific-python/cookie
     rev: 2025.01.22
     hooks:

pyproject.toml

@@ -356,6 +356,7 @@ enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
 [[tool.mypy.overrides]]
 module = [
     "tests.package_with_entrypoint.*",
+    "zarr.testing.stateful",
     "tests.test_codecs.test_transpose",
     "tests.test_config"
 ]
@@ -365,7 +366,6 @@ strict = false
 # and fix the errors
 [[tool.mypy.overrides]]
 module = [
-    "zarr.testing.stateful",  # lots of hypothesis decorator errors
     "tests.test_codecs.test_codecs",
     "tests.test_metadata.*",
     "tests.test_store.*",

src/zarr/core/array.py

@@ -672,7 +672,7 @@ def _create_metadata_v3(
         fill_value: Any | None = None,
         chunk_key_encoding: ChunkKeyEncodingLike | None = None,
         codecs: Iterable[Codec | dict[str, JSON]] | None = None,
-        dimension_names: Iterable[str] | None = None,
+        dimension_names: Iterable[str | None] | None = None,
         attributes: dict[str, JSON] | None = None,
     ) -> ArrayV3Metadata:
         """
@@ -3821,7 +3821,7 @@ async def from_array(
     zarr_format: ZarrFormat | None = None,
     attributes: dict[str, JSON] | None = None,
     chunk_key_encoding: ChunkKeyEncodingLike | None = None,
-    dimension_names: Iterable[str] | None = None,
+    dimension_names: Iterable[str | None] | None = None,
     storage_options: dict[str, Any] | None = None,
     overwrite: bool = False,
     config: ArrayConfig | ArrayConfigLike | None = None,
@@ -4083,7 +4083,7 @@ async def init_array(
     zarr_format: ZarrFormat | None = 3,
     attributes: dict[str, JSON] | None = None,
     chunk_key_encoding: ChunkKeyEncodingLike | None = None,
-    dimension_names: Iterable[str] | None = None,
+    dimension_names: Iterable[str | None] | None = None,
     overwrite: bool = False,
     config: ArrayConfigLike | None,
 ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]:
@@ -4477,7 +4477,7 @@ def _parse_keep_array_attr(
     order: MemoryOrder | None,
     zarr_format: ZarrFormat | None,
     chunk_key_encoding: ChunkKeyEncodingLike | None,
-    dimension_names: Iterable[str] | None,
+    dimension_names: Iterable[str | None] | None,
 ) -> tuple[
     ChunkCoords | Literal["auto"],
     ShardsLike | None,
@@ -4488,7 +4488,7 @@ def _parse_keep_array_attr(
     MemoryOrder | None,
     ZarrFormat,
     ChunkKeyEncodingLike | None,
-    Iterable[str] | None,
+    Iterable[str | None] | None,
 ]:
     if isinstance(data, Array):
         if chunks == "keep":

src/zarr/core/metadata/v3.py

@@ -242,7 +242,7 @@ class ArrayV3Metadata(Metadata):
     fill_value: Any
     codecs: tuple[Codec, ...]
     attributes: dict[str, Any] = field(default_factory=dict)
-    dimension_names: tuple[str, ...] | None = None
+    dimension_names: tuple[str | None, ...] | None = None
     zarr_format: Literal[3] = field(default=3, init=False)
     node_type: Literal["array"] = field(default="array", init=False)
     storage_transformers: tuple[dict[str, JSON], ...]
@@ -257,7 +257,7 @@ def __init__(
         fill_value: Any,
         codecs: Iterable[Codec | dict[str, JSON]],
         attributes: dict[str, JSON] | None,
-        dimension_names: Iterable[str] | None,
+        dimension_names: Iterable[str | None] | None,
         storage_transformers: Iterable[dict[str, JSON]] | None = None,
     ) -> None:
         """

src/zarr/testing/stateful.py

@@ -326,8 +326,8 @@ def init_store(self) -> None:
         self.store.clear()
 
     @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}")
+    def set(self, key: str, data: bytes) -> None:
+        note(f"(set) Setting {key!r} with {data!r}")
         assert not self.store.read_only
         data_buf = cpu.Buffer.from_bytes(data)
         self.store.set(key, data_buf)

src/zarr/testing/strategies.py

@@ -1,10 +1,12 @@
 import math
 import sys
+from collections.abc import Callable, Mapping
 from typing import Any, Literal
 
 import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
 import numpy as np
+import numpy.typing as npt
 from hypothesis import event
 from hypothesis.strategies import SearchStrategy
 
@@ -14,7 +16,7 @@
 from zarr.core.array import Array
 from zarr.core.chunk_grids import RegularChunkGrid
 from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding
-from zarr.core.common import ZarrFormat
+from zarr.core.common import JSON, ZarrFormat
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
 from zarr.core.sync import sync
 from zarr.storage import MemoryStore, StoreLike
@@ -30,17 +32,17 @@
 )
 
 
-@st.composite  # type: ignore[misc]
-def keys(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> Any:
+@st.composite
+def keys(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str:
     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:
+@st.composite
+def paths(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str:
     return draw(st.just("/") | keys(max_num_nodes=max_num_nodes))
 
 
-def v3_dtypes() -> st.SearchStrategy[np.dtype]:
+def v3_dtypes() -> st.SearchStrategy[np.dtype[Any]]:
     return (
         npst.boolean_dtypes()
         | npst.integer_dtypes(endianness="=")
@@ -54,7 +56,7 @@ def v3_dtypes() -> st.SearchStrategy[np.dtype]:
     )
 
 
-def v2_dtypes() -> st.SearchStrategy[np.dtype]:
+def v2_dtypes() -> st.SearchStrategy[np.dtype[Any]]:
     return (
         npst.boolean_dtypes()
         | npst.integer_dtypes(endianness="=")
@@ -107,7 +109,9 @@ def clear_store(x: Store) -> Store:
     .filter(lambda name: name.lower() != "zarr.json")
 )
 array_names = node_names
-attrs = st.none() | st.dictionaries(_attr_keys, _attr_values)
+attrs: st.SearchStrategy[Mapping[str, JSON] | None] = st.none() | st.dictionaries(
+    _attr_keys, _attr_values
+)
 # 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.
@@ -118,19 +122,19 @@ def clear_store(x: Store) -> Store:
 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]
+@st.composite
 def dimension_names(draw: st.DrawFn, *, ndim: int | None = None) -> list[None | str] | None:
     simple_text = st.text(zarr_key_chars, min_size=0)
-    return draw(st.none() | st.lists(st.none() | simple_text, min_size=ndim, max_size=ndim))  # type: ignore[no-any-return]
+    return draw(st.none() | st.lists(st.none() | simple_text, min_size=ndim, max_size=ndim))  # type: ignore[arg-type]
 
 
-@st.composite  # type: ignore[misc]
+@st.composite
 def array_metadata(
     draw: st.DrawFn,
     *,
-    array_shapes: st.SearchStrategy[tuple[int, ...]] = npst.array_shapes,
+    array_shapes: Callable[..., st.SearchStrategy[tuple[int, ...]]] = npst.array_shapes,
     zarr_formats: st.SearchStrategy[Literal[2, 3]] = zarr_formats,
-    attributes: st.SearchStrategy[dict[str, Any]] = attrs,
+    attributes: SearchStrategy[Mapping[str, JSON] | None] = attrs,
 ) -> ArrayV2Metadata | ArrayV3Metadata:
     zarr_format = draw(zarr_formats)
     # separator = draw(st.sampled_from(['/', '\\']))
@@ -146,7 +150,7 @@ def array_metadata(
             dtype=dtype,
             fill_value=fill_value,
             order=draw(st.sampled_from(["C", "F"])),
-            attributes=draw(attributes),
+            attributes=draw(attributes),  # type: ignore[arg-type]
             dimension_separator=draw(st.sampled_from([".", "/"])),
             filters=None,
             compressor=None,
@@ -157,22 +161,22 @@ def array_metadata(
             data_type=dtype,
             chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape),
             fill_value=fill_value,
-            attributes=draw(attributes),
+            attributes=draw(attributes),  # type: ignore[arg-type]
             dimension_names=draw(dimension_names(ndim=ndim)),
             chunk_key_encoding=DefaultChunkKeyEncoding(separator="/"),  # FIXME
             codecs=[BytesCodec()],
             storage_transformers=(),
         )
 
 
-@st.composite  # type: ignore[misc]
+@st.composite
 def numpy_arrays(
     draw: st.DrawFn,
     *,
     shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes,
     dtype: np.dtype[Any] | None = None,
-    zarr_formats: st.SearchStrategy[ZarrFormat] | None = zarr_formats,
-) -> Any:
+    zarr_formats: st.SearchStrategy[ZarrFormat] = zarr_formats,
+) -> npt.NDArray[Any]:
     """
     Generate numpy arrays that can be saved in the provided Zarr format.
     """
@@ -186,7 +190,7 @@ def numpy_arrays(
     return draw(npst.arrays(dtype=dtype, shape=shapes))
 
 
-@st.composite  # type: ignore[misc]
+@st.composite
 def chunk_shapes(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[int, ...]:
     # We want this strategy to shrink towards arrays with smaller number of chunks
     # 1. st.integers() shrinks towards smaller values. So we use that to generate number of chunks
@@ -208,7 +212,7 @@ def chunk_shapes(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[int, ...]:
     return chunks
 
 
-@st.composite  # type: ignore[misc]
+@st.composite
 def shard_shapes(
     draw: st.DrawFn, *, shape: tuple[int, ...], chunk_shape: tuple[int, ...]
 ) -> tuple[int, ...]:
@@ -220,9 +224,11 @@ def shard_shapes(
     return tuple(m * c for m, c in zip(multiples, chunk_shape, strict=True))
 
 
-@st.composite  # type: ignore[misc]
+@st.composite
 def np_array_and_chunks(
-    draw: st.DrawFn, *, arrays: st.SearchStrategy[np.ndarray] = numpy_arrays
+    draw: st.DrawFn,
+    *,
+    arrays: st.SearchStrategy[npt.NDArray[Any]] = numpy_arrays(),  # noqa: B008
 ) -> tuple[np.ndarray, tuple[int, ...]]:  # type: ignore[type-arg]
     """A hypothesis strategy to generate small sized random arrays.
 
@@ -232,14 +238,14 @@ def np_array_and_chunks(
     return (array, draw(chunk_shapes(shape=array.shape)))
 
 
-@st.composite  # type: ignore[misc]
+@st.composite
 def arrays(
     draw: st.DrawFn,
     *,
     shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes,
     compressors: st.SearchStrategy = compressors,
     stores: st.SearchStrategy[StoreLike] = stores,
-    paths: st.SearchStrategy[str | None] = paths(),  # noqa: B008
+    paths: st.SearchStrategy[str] = paths(),  # noqa: B008
     array_names: st.SearchStrategy = array_names,
     arrays: st.SearchStrategy | None = None,
     attrs: st.SearchStrategy = attrs,
@@ -296,7 +302,7 @@ def arrays(
     return a
 
 
-@st.composite  # type: ignore[misc]
+@st.composite
 def simple_arrays(
     draw: st.DrawFn,
     *,
@@ -317,7 +323,7 @@ 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]
+@st.composite
 def end_slices(draw: st.DrawFn, *, shape: tuple[int]) -> Any:
     """
     A strategy that slices ranges that include the last chunk.
@@ -332,23 +338,32 @@ def end_slices(draw: st.DrawFn, *, shape: tuple[int]) -> Any:
     return tuple(slicers)
 
 
-@st.composite  # type: ignore[misc]
-def basic_indices(draw: st.DrawFn, *, shape: tuple[int], **kwargs: Any) -> Any:
+@st.composite
+def basic_indices(
+    draw: st.DrawFn,
+    *,
+    shape: tuple[int],
+    min_dims: int = 0,
+    max_dims: int | None = None,
+    allow_newaxis: bool = False,
+    allow_ellipsis: bool = True,
+) -> Any:
     """Basic indices without unsupported negative slices."""
-    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))
-            )
-        )
+    strategy = npst.basic_indices(
+        shape=shape,
+        min_dims=min_dims,
+        max_dims=max_dims,
+        allow_newaxis=allow_newaxis,
+        allow_ellipsis=allow_ellipsis,
+    ).filter(
+        lambda idxr: (not (is_negative_slice(idxr) or 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]
+@st.composite
 def orthogonal_indices(
     draw: st.DrawFn, *, shape: tuple[int]
 ) -> tuple[tuple[np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]:
@@ -386,8 +401,8 @@ def orthogonal_indices(
 
 
 def key_ranges(
-    keys: SearchStrategy = node_names, max_size: int = sys.maxsize
-) -> SearchStrategy[list[int]]:
+    keys: SearchStrategy[str] = node_names, max_size: int = sys.maxsize
+) -> SearchStrategy[list[tuple[str, RangeByteRequest]]]:
     """
     Function to generate key_ranges strategy for get_partial_values()
     returns list strategy w/ form::