feat(zarr-metadata): add Hypothesis strategies for metadata fields

packages/zarr-metadata/changes/414.feature.md

@@ -0,0 +1 @@
+Added `zarr_metadata.testing.strategies`, Hypothesis strategies that draw valid Zarr v3 metadata fields, behind a new `testing` extra (`pip install zarr-metadata[testing]`). There is one strategy for every codec, data type, chunk grid and chunk key encoding the package defines that holds no other field. `fields` draws every envelope spelling the spec allows, or only the canonical one with `spellings="canonical"`. Codecs can be drawn for the `Chunk` they are handed, chunk grids for an array shape, and fill values for a data type. The strategies are a mapping keyed by definition, so an extension author can add strategies for their own definitions.

packages/zarr-metadata/docs/api/testing.md

@@ -0,0 +1,8 @@
+---
+title: testing.strategies
+---
+
+::: zarr_metadata.testing.strategies
+    options:
+      inherited_members: false
+      show_if_no_docstring: false

packages/zarr-metadata/mkdocs.yml

@@ -16,6 +16,7 @@ nav:
       - api/index.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_metadata.model</code>': api/model.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_metadata.pydantic</code>': api/pydantic.md
+      - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_metadata.testing.strategies</code>': api/testing.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_metadata.typed_json</code>': api/typed_json.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_metadata.v2</code>': api/v2.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_metadata.v3</code>':

packages/zarr-metadata/pyproject.toml

@@ -39,6 +39,11 @@ dependencies = [
     "typing_extensions>=4.16",
 ]
 
+[project.optional-dependencies]
+# `zarr_metadata.testing.strategies`. 6.85: first release with
+# `characters(exclude_categories=...)`.
+testing = ["hypothesis>=6.85"]
+
 [project.urls]
 Homepage = "https://github.com/zarr-developers/zarr-python"
 Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-metadata"

packages/zarr-metadata/src/zarr_metadata/testing/__init__.py

@@ -0,0 +1,6 @@
+"""Testing utilities for code that reads or writes Zarr metadata.
+
+`zarr_metadata.testing.strategies` holds Hypothesis strategies that draw
+valid metadata. It needs the `testing` extra:
+`pip install zarr-metadata[testing]`.
+"""

packages/zarr-metadata/src/zarr_metadata/testing/_chunk_grid.py

@@ -0,0 +1,107 @@
+"""Strategies for each chunk grid this package defines.
+
+Over an array of a known shape, a grid is drawn to fit it: a chunk
+length for each dimension, never zero ("Chunk sizes must be greater
+than zero", https://github.com/zarr-developers/zarr-specs/blob/fc7dd9c9beb5a50b87f9b08b00bf50fc0048482f/docs/v3/chunk-grids/regular-grid/index.rst#L40),
+and rectilinear edges that cover each dimension. Over no shape, what is
+drawn is what the grid's rules accept.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Final
+
+from hypothesis import strategies as st
+
+from zarr_metadata.testing._strategy import ChunkGridStrategy
+from zarr_metadata.v3.chunk_grid.rectilinear import RECTILINEAR_CHUNK_GRID
+from zarr_metadata.v3.chunk_grid.regular import REGULAR_CHUNK_GRID
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from zarr_metadata._common import JSONValue
+    from zarr_metadata.v3._definition import ChunkGridDefinition
+
+_MAX_RANK: Final = 5
+"""The most dimensions a grid is drawn with, over no shape."""
+
+_MAX_LENGTH: Final = 16
+"""The longest chunk drawn along a dimension whose length is unknown."""
+
+_MAX_RUNS: Final = 5
+"""The most runs of equal lengths a rectilinear dimension is drawn with."""
+
+_MAX_COUNT: Final = 4
+"""The most chunks a run is drawn with."""
+
+
+def _length(extent: int | None) -> st.SearchStrategy[int]:
+    """A chunk length along a dimension of `extent`: at least 1, and up to twice the extent when it is known and not 0."""
+    if extent is None or extent == 0:
+        return st.integers(1, _MAX_LENGTH)
+    return st.integers(1, 2 * extent)
+
+
+def _regular(shape: tuple[int, ...] | None) -> st.SearchStrategy[Mapping[str, JSONValue]]:
+    """A chunk length for each dimension."""
+    if shape is None:
+        lengths = st.lists(_length(None), max_size=_MAX_RANK).map(tuple)
+    else:
+        lengths = st.tuples(*(_length(extent) for extent in shape))
+    return lengths.map(lambda chunk_shape: {"chunk_shape": chunk_shape})
+
+
+def _spelled_run(draw: st.DrawFn, size: int, count: int) -> list[JSONValue]:
+    """`count` chunks of `size`: one `[size, count]` pair, or the size written once for each."""
+    if draw(st.booleans()):
+        return [(size, count)]
+    return [size] * count
+
+
+@st.composite
+def _edges(draw: st.DrawFn, extent: int | None) -> JSONValue:
+    """Edge lengths along a dimension, covering `extent`, in runs spelled as pairs or written out.
+
+    Runs are drawn apart, so two of one size can sit side by side and a
+    pair can count one chunk: spellings the canonical form merges. A
+    dimension of length zero is covered by no edges at all.
+    """
+    min_runs = 0 if extent == 0 else 1
+    runs = draw(
+        st.lists(
+            st.tuples(st.integers(1, _MAX_LENGTH), st.integers(1, _MAX_COUNT)),
+            min_size=min_runs,
+            max_size=_MAX_RUNS,
+        )
+    )
+    covered = sum(size * count for size, count in runs)
+    if extent is not None and covered < extent:
+        runs.append((extent - covered, 1))
+    spelled: list[JSONValue] = []
+    for size, count in runs:
+        spelled.extend(_spelled_run(draw, size, count))
+    return tuple(spelled)
+
+
+def _rectilinear_dimension(extent: int | None) -> st.SearchStrategy[JSONValue]:
+    """A bare length, which repeats until it covers the dimension, or edges that cover it."""
+    return _length(extent) | _edges(extent)
+
+
+def _rectilinear(shape: tuple[int, ...] | None) -> st.SearchStrategy[Mapping[str, JSONValue]]:
+    """Inline chunk shapes, one entry for each dimension."""
+    if shape is None:
+        dimensions = st.lists(_rectilinear_dimension(None), max_size=_MAX_RANK).map(tuple)
+    else:
+        dimensions = st.tuples(*(_rectilinear_dimension(extent) for extent in shape))
+    return dimensions.map(lambda chunk_shapes: {"kind": "inline", "chunk_shapes": chunk_shapes})
+
+
+CHUNK_GRID_STRATEGIES: Final[Mapping[ChunkGridDefinition[Any], ChunkGridStrategy]] = {
+    REGULAR_CHUNK_GRID: ChunkGridStrategy(configurations=_regular),
+    RECTILINEAR_CHUNK_GRID: ChunkGridStrategy(configurations=_rectilinear),
+}
+"""A strategy for each chunk grid."""
+
+__all__ = ["CHUNK_GRID_STRATEGIES"]

packages/zarr-metadata/src/zarr_metadata/testing/_chunk_key_encoding.py

@@ -0,0 +1,40 @@
+"""Strategies for each chunk key encoding this package defines."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Final
+
+from hypothesis import strategies as st
+
+from zarr_metadata.testing._strategy import ChunkKeyEncodingStrategy
+from zarr_metadata.v3.chunk_key_encoding.default import (
+    DEFAULT_CHUNK_KEY_ENCODING,
+    DEFAULT_CHUNK_KEY_ENCODING_SEPARATOR,
+)
+from zarr_metadata.v3.chunk_key_encoding.v2 import (
+    V2_CHUNK_KEY_ENCODING,
+    V2_CHUNK_KEY_ENCODING_SEPARATOR,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from zarr_metadata.v3._definition import ChunkKeyEncodingDefinition
+
+CHUNK_KEY_ENCODING_STRATEGIES: Final[
+    Mapping[ChunkKeyEncodingDefinition[Any], ChunkKeyEncodingStrategy]
+] = {
+    DEFAULT_CHUNK_KEY_ENCODING: ChunkKeyEncodingStrategy(
+        configurations=st.fixed_dictionaries(
+            {}, optional={"separator": st.sampled_from(DEFAULT_CHUNK_KEY_ENCODING_SEPARATOR)}
+        )
+    ),
+    V2_CHUNK_KEY_ENCODING: ChunkKeyEncodingStrategy(
+        configurations=st.fixed_dictionaries(
+            {}, optional={"separator": st.sampled_from(V2_CHUNK_KEY_ENCODING_SEPARATOR)}
+        )
+    ),
+}
+"""A strategy for each chunk key encoding: a `separator`, or none, which means the encoding's default."""
+
+__all__ = ["CHUNK_KEY_ENCODING_STRATEGIES"]

packages/zarr-metadata/src/zarr_metadata/testing/_codec.py

@@ -0,0 +1,136 @@
+"""Strategies for each codec this package defines that holds no other field.
+
+A codec handed an array draws what fits the chunk: a `bytes` codec an
+`endian` for a type whose values take several bytes, a `transpose` an
+`order` with an entry for each axis. Handed `Chunk()`, nothing is known
+of the chunk, and what is drawn is what the codec's rules accept.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Final
+
+from hypothesis import strategies as st
+
+from zarr_metadata.testing._data_type import FLOAT_WIDTHS, INTEGER_RANGES, finite_floats
+from zarr_metadata.testing._strategy import CodecStrategy
+from zarr_metadata.v3._definition import Chunk, storage_of
+from zarr_metadata.v3.codec.blosc import BLOSC_CNAME, BLOSC_CODEC, BLOSC_NO_SHUFFLE, BLOSC_SHUFFLE
+from zarr_metadata.v3.codec.bytes import BYTES_CODEC, ENDIANNESS
+from zarr_metadata.v3.codec.crc32c import CRC32C_CODEC
+from zarr_metadata.v3.codec.gzip import GZIP_CODEC
+from zarr_metadata.v3.codec.scale_offset import SCALE_OFFSET_CODEC
+from zarr_metadata.v3.codec.transpose import TRANSPOSE_CODEC
+from zarr_metadata.v3.codec.zstd import ZSTD_CODEC, ZSTD_MAX_LEVEL, ZSTD_MIN_LEVEL
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from zarr_metadata._common import JSONValue
+    from zarr_metadata.v3._definition import CodecDefinition
+
+_MAX_RANK: Final = 5
+"""The most axes an `order` is drawn with, when the chunk does not say."""
+
+_BLOSC_MAX_TYPESIZE: Final = 255
+"""The widest element c-blosc takes: `BLOSC_MAX_TYPESIZE`."""
+
+_BLOSC_MAX_BLOCKSIZE: Final = 2**31 - 1
+"""The largest block c-blosc takes, a signed 32-bit size; 0 lets it choose."""
+
+_ENDIAN: Final = st.sampled_from(ENDIANNESS)
+
+
+def _bytes(chunk: Chunk) -> st.SearchStrategy[Mapping[str, JSONValue]] | None:
+    """An `endian`, required for a data type whose values take several bytes; None for one whose values vary in size."""
+    storage = None if chunk.data_type is None else storage_of(chunk.data_type)
+    if storage == "variable_length":
+        return None
+    if storage == "multi_byte":
+        return st.fixed_dictionaries({"endian": _ENDIAN})
+    return st.fixed_dictionaries({}, optional={"endian": _ENDIAN})
+
+
+@st.composite
+def _blosc(draw: st.DrawFn) -> Mapping[str, JSONValue]:
+    """A compressor, a level, a shuffle and a block size; and a `typesize`, which only `noshuffle` may leave out.
+
+    Under `noshuffle` the spec says `typesize` "is ignored", so it is
+    drawn there too, a spelling that is not the simplest.
+    """
+    shuffle = draw(st.sampled_from(BLOSC_SHUFFLE))
+    typesize = st.integers(1, _BLOSC_MAX_TYPESIZE)
+    required: dict[str, st.SearchStrategy[JSONValue]] = {
+        "cname": st.sampled_from(BLOSC_CNAME),
+        "clevel": st.integers(0, 9),
+        "shuffle": st.just(shuffle),
+        "blocksize": st.integers(0, _BLOSC_MAX_BLOCKSIZE),
+    }
+    if shuffle == BLOSC_NO_SHUFFLE:
+        return draw(st.fixed_dictionaries(required, optional={"typesize": typesize}))
+    return draw(st.fixed_dictionaries({**required, "typesize": typesize}))
+
+
+def _order(rank: int) -> st.SearchStrategy[Mapping[str, JSONValue]]:
+    return st.permutations(range(rank)).map(lambda order: {"order": tuple(order)})
+
+
+def _transpose(chunk: Chunk) -> st.SearchStrategy[Mapping[str, JSONValue]]:
+    """A permutation of the chunk's axes; of some number of axes, when the chunk does not say."""
+    if chunk.rank is not None:
+        return _order(chunk.rank)
+    return st.integers(0, _MAX_RANK).flatmap(_order)
+
+
+def _scale_offset(chunk: Chunk) -> st.SearchStrategy[Mapping[str, JSONValue]] | None:
+    """An `offset` and a `scale`, each a finite value of the data type the codec is handed, the scale not zero.
+
+    Each is a fill value of the chunk's data type
+    (https://github.com/zarr-developers/zarr-extensions/blob/4da7b37a84f76e660902f6d3de3eaef0e0febae6/codecs/scale_offset/README.md?plain=1#L39-L43),
+    drawn finite, and the scale non-zero, so that decoding, which divides
+    by it, is defined. Handed a data type the codec's registry entry does
+    not list -- truth values, raw bits, complex numbers, numpy times --
+    nothing is drawn: for some the spec refuses it, and for the others
+    leaves it open. Handed a chunk whose data type is unknown, the values
+    are small non-negative integers, which every listed type holds.
+    """
+    data_type = chunk.data_type
+    read = data_type is not None and data_type.resolution == "read"
+    definition = data_type.definition if data_type is not None and read else None
+    if definition is None:
+        offset: st.SearchStrategy[JSONValue] = st.integers(0, 127)
+        scale: st.SearchStrategy[JSONValue] = st.integers(1, 127)
+    elif definition in INTEGER_RANGES:
+        low, high = INTEGER_RANGES[definition]
+        offset = st.integers(low, high)
+        scale = st.integers(low, high).filter(lambda value: value != 0)
+    elif definition in FLOAT_WIDTHS:
+        offset = finite_floats(FLOAT_WIDTHS[definition])
+        scale = finite_floats(FLOAT_WIDTHS[definition]).filter(lambda value: value != 0)
+    else:
+        return None
+    return st.fixed_dictionaries({}, optional={"offset": offset, "scale": scale})
+
+
+def _unchunked(configurations: st.SearchStrategy[Mapping[str, JSONValue]]) -> CodecStrategy:
+    """A codec whose configurations do not depend on the chunk it is handed."""
+    return CodecStrategy(configurations=lambda _chunk: configurations)
+
+
+CODEC_STRATEGIES: Final[Mapping[CodecDefinition[Any], CodecStrategy]] = {
+    BLOSC_CODEC: _unchunked(_blosc()),
+    BYTES_CODEC: CodecStrategy(configurations=_bytes),
+    CRC32C_CODEC: _unchunked(st.just({})),
+    GZIP_CODEC: _unchunked(st.fixed_dictionaries({"level": st.integers(0, 9)})),
+    SCALE_OFFSET_CODEC: CodecStrategy(configurations=_scale_offset),
+    TRANSPOSE_CODEC: CodecStrategy(configurations=_transpose),
+    ZSTD_CODEC: _unchunked(
+        st.fixed_dictionaries(
+            {"level": st.integers(ZSTD_MIN_LEVEL, ZSTD_MAX_LEVEL)},
+            optional={"checksum": st.booleans()},
+        )
+    ),
+}
+"""A strategy for each codec that holds no other field."""
+
+__all__ = ["CODEC_STRATEGIES"]