8000 Add Array metadata strategy (#2813) · tomwhite/zarr-python@99621ec · GitHub
[go: up one dir, main page]
Skip to content

Commit 99621ec

Browse files
dcheriandcherian
authored
Add Array metadata strategy (zarr-developers#2813)
* Add array_metadata strategy * Fix merge
1 parent 870265a commit 99621ec

File tree

3 files changed

+77
-3
lines changed

3 files changed

+77
-3
lines changed

‎changes/2813.feature.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add `zarr.testing.strategies.array_metadata` to generate ArrayV2Metadata and ArrayV3Metadata instances.

‎src/zarr/testing/strategies.py

Lines changed: 59 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import sys
2-
from typing import Any
2+
from typing import Any, Literal
33

44
import hypothesis.extra.numpy as npst
55
import hypothesis.strategies as st
@@ -8,9 +8,13 @@
88
from hypothesis.strategies import SearchStrategy
99

1010
import zarr
11-
from zarr.abc.store import RangeByteRequest
11+
from zarr.abc.store import RangeByteRequest, Store
12+
from zarr.codecs.bytes import BytesCodec
1213
from zarr.core.array import Array
14+
from zarr.core.chunk_grids import RegularChunkGrid
15+
from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding
1316
from zarr.core.common import ZarrFormat
17+
from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
1418
from zarr.core.sync import sync
1519
from zarr.storage import MemoryStore, StoreLike
1620
from zarr.storage._common import _dereference_path
@@ -67,6 +71,11 @@ def safe_unicode_for_dtype(dtype: np.dtype[np.str_]) -> st.SearchStrategy[str]:
6771
)
6872

6973

74+
def clear_store(x: Store) -> Store:
75+
sync(x.clear())
76+
return x
77+
78+
7079
# From https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#node-names
7180
# 1. must not be the empty string ("")
7281
# 2. must not include the character "/"
@@ -85,12 +94,59 @@ def safe_unicode_for_dtype(dtype: np.dtype[np.str_]) -> st.SearchStrategy[str]:
8594
# st.builds will only call a new store constructor for different keyword arguments
8695
# i.e. stores.examples() will always return the same object per Store class.
8796
# So we map a clear to reset the store.
88-
stores = st.builds(MemoryStore, st.just({})).map(lambda x: sync(x.clear()))
97+
stores = st.builds(MemoryStore, st.just({})).map(clear_store)
8998
compressors = st.sampled_from([None, "default"])
9099
zarr_formats: st.SearchStrategy[ZarrFormat] = st.sampled_from([2, 3])
91100
array_shapes = npst.array_shapes(max_dims=4, min_side=0)
92101

93102

103+
@st.composite # type: ignore[misc]
104+
def dimension_names(draw: st.DrawFn, *, ndim: int | None = None) -> list[None | str] | None:
105+
simple_text = st.text(zarr_key_chars, min_size=0)
106+
return draw(st.none() | st.lists(st.none() | simple_text, min_size=ndim, max_size=ndim)) # type: ignore[no-any-return]
107+
108+
109+
@st.composite # type: ignore[misc]
110+
def array_metadata(
111+
draw: st.DrawFn,
112+
*,
113+
array_shapes: st.SearchStrategy[tuple[int, ...]] = npst.array_shapes,
114+
zarr_formats: st.SearchStrategy[Literal[2, 3]] = zarr_formats,
115+
attributes: st.SearchStrategy[dict[str, Any]] = attrs,
116+
) -> ArrayV2Metadata | ArrayV3Metadata:
117+
zarr_format = draw(zarr_formats)
118+
# separator = draw(st.sampled_from(['/', '\\']))
119+
shape = draw(array_shapes())
120+
ndim = len(shape)
121+
chunk_shape = draw(array_shapes(min_dims=ndim, max_dims=ndim))
122+
dtype = draw(v3_dtypes())
123+
fill_value = draw(npst.from_dtype(dtype))
124+
if zarr_format == 2:
125+
return ArrayV2Metadata(
126+
shape=shape,
127+
chunks=chunk_shape,
128+
dtype=dtype,
129+
fill_value=fill_value,
130+
order=draw(st.sampled_from(["C", "F"])),
131+
attributes=draw(attributes),
132+
dimension_separator=draw(st.sampled_from([".", "/"])),
133+
filters=None,
134+
compressor=None,
135+
)
136+
else:
137+
return ArrayV3Metadata(
138+
shape=shape,
139+
data_type=dtype,
140+
chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape),
141+
fill_value=fill_value,
142+
attributes=draw(attributes),
143+
dimension_names=draw(dimension_names(ndim=ndim)),
144+
chunk_key_encoding=DefaultChunkKeyEncoding(separator="/"), # FIXME
145+
codecs=[BytesCodec()],
146+
storage_transformers=(),
147+
)
148+
149+
94150
@st.composite # type: ignore[misc]
95151
def numpy_arrays(
96152
draw: st.DrawFn,

‎tests/test_properties.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,23 @@
22
import pytest
33
from numpy.testing import assert_array_equal
44

5+
from zarr.core.buffer import default_buffer_prototype
6+
57
pytest.importorskip("hypothesis")
68

79
import hypothesis.extra.numpy as npst
810
import hypothesis.strategies as st
911
from hypothesis import given
1012

13+
from zarr.abc.store import Store
14+
from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
1115
from zarr.testing.strategies import (
16+
array_metadata,
1217
arrays,
1318
basic_indices,
1419
numpy_arrays,
1520
orthogonal_indices,
21+
stores,
1622
zarr_formats,
1723
)
1824

@@ -64,6 +70,17 @@ def test_vindex(data: st.DataObject) -> None:
6470
assert_array_equal(nparray[indexer], actual)
6571

6672

73+
@given(store=stores, meta=array_metadata()) # type: ignore[misc]
74+
async def test_roundtrip_array_metadata(
75+
store: Store, meta: ArrayV2Metadata | ArrayV3Metadata
76+
) -> None:
77+
asdict = meta.to_buffer_dict(prototype=default_buffer_prototype())
78+
for key, expected in asdict.items():
79+
await store.set(f"0/{key}", expected)
80+
actual = await store.get(f"0/{key}", prototype=default_buffer_prototype())
81+
assert actual == expected
82+
83+
6784
# @st.composite
6885
# def advanced_indices(draw, *, shape):
6986
# basic_idxr = draw(

0 commit comments

Comments
 (0)
0