Fix fill_value serialization of NaN; add property-based tests

Author: moradology

src/zarr/core/metadata/v2.py

@@ -19,6 +19,7 @@
     from zarr.core.common import ChunkCoords
 
 import json
+import numbers
 from dataclasses import dataclass, field, fields, replace
 
 import numcodecs
@@ -149,18 +150,20 @@ def _json_convert(
         json_indent = config.get("json_indent")
         return {
             ZARRAY_JSON: prototype.buffer.from_bytes(
-                json.dumps(zarray_dict, default=_json_convert, indent=json_indent).encode()
+                json.dumps(
+                    zarray_dict, default=_json_convert, indent=json_indent, allow_nan=False
+                ).encode()
             ),
             ZATTRS_JSON: prototype.buffer.from_bytes(
-                json.dumps(zattrs_dict, indent=json_indent).encode()
+                json.dumps(zattrs_dict, indent=json_indent, allow_nan=False).encode()
             ),
         }
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> ArrayV2Metadata:
-        # make a copy to protect the original from modification
+        # Make a copy to protect the original from modification.
         _data = data.copy()
-        # check that the zarr_format attribute is correct
+        # Check that the zarr_format attribute is correct.
         _ = parse_zarr_format(_data.pop("zarr_format"))
         dtype = parse_dtype(_data["dtype"])
 
@@ -169,20 +172,46 @@ def from_dict(cls, data: dict[str, Any]) -> ArrayV2Metadata:
             if fill_value_encoded is not None:
                 fill_value = base64.standard_b64decode(fill_value_encoded)
                 _data["fill_value"] = fill_value
-
-        # zarr v2 allowed arbitrary keys here.
-        # We don't want the ArrayV2Metadata constructor to fail just because someone put an
-        # extra key in the metadata.
+        else:
+            fill_value = _data.get("fill_value")
+            if fill_value is not None:
+                if np.issubdtype(dtype, np.datetime64):
+                    if fill_value == "NaT":
+                        _data["fill_value"] = np.array("NaT", dtype=dtype)[()]
+                    else:
+                        _data["fill_value"] = np.array(fill_value, dtype=dtype)[()]
+                elif dtype.kind == "c" and isinstance(fill_value, list):
+                    if len(fill_value) == 2:
+                        val = complex(float(fill_value[0]), float(fill_value[1]))
+                        _data["fill_value"] = np.array(val, dtype=dtype)[()]
+                elif dtype.kind in "f" and isinstance(fill_value, str):
+                    if fill_value in {"NaN", "Infinity", "-Infinity"}:
+                        _data["fill_value"] = np.array(fill_value, dtype=dtype)[()]
+        # zarr v2 allowed arbitrary keys in the metadata.
+        # Filter the keys to only those expected by the constructor.
         expected = {x.name for x in fields(cls)}
-        # https://github.com/zarr-developers/zarr-python/issues/2269
-        # handle the renames
         expected |= {"dtype", "chunks"}
-
         _data = {k: v for k, v in _data.items() if k in expected}
 
         return cls(**_data)
 
     def to_dict(self) -> dict[str, JSON]:
+        def _sanitize_fill_value(fv: Any):
+            if fv is None:
+                return fv
+            elif isinstance(fv, np.datetime64):
+                if np.isnat(fv):
+                    return "NaT"
+                return np.datetime_as_string(fv)
+            elif isinstance(fv, numbers.Real):
+                if np.isnan(fv):
+                    fv = "NaN"
+                elif np.isinf(fv):
+                    fv = "Infinity" if fv > 0 else "-Infinity"
+            elif isinstance(fv, numbers.Complex):
+                fv = [_sanitize_fill_value(fv.real), _sanitize_fill_value(fv.imag)]
+            return fv
+
         zarray_dict = super().to_dict()
 
         if self.dtype.kind in "SV" and self.fill_value is not None:
@@ -192,6 +221,7 @@ def to_dict(self) -> dict[str, JSON]:
             fill_value = base64.standard_b64encode(cast(bytes, self.fill_value)).decode("ascii")
             zarray_dict["fill_value"] = fill_value
 
+        zarray_dict["fill_value"] = _sanitize_fill_value(zarray_dict["fill_value"])
         _ = zarray_dict.pop("dtype")
         dtype_json: JSON
         # In the case of zarr v2, the simplest i.e., '|VXX' dtype is represented as a string
@@ -300,7 +330,6 @@ def parse_fill_value(fill_value: object, dtype: np.dtype[Any]) -> Any:
     -------
         An instance of `dtype`, or `None`, or any python object (in the case of an object dtype)
     """
-
     if fill_value is None or dtype.hasobject:
         # no fill value
         pass

src/zarr/testing/stateful.py

@@ -85,7 +85,7 @@ def add_group(self, name: str, data: DataObject) -> None:
     @rule(
         data=st.data(),
         name=node_names,
-        array_and_chunks=np_array_and_chunks(arrays=numpy_arrays(zarr_formats=st.just(3))),
+        array_and_chunks=np_array_and_chunks(nparrays=numpy_arrays(zarr_formats=st.just(3))),
     )
     def add_array(
         self,

src/zarr/testing/strategies.py

@@ -3,8 +3,9 @@
 
 import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
+import numcodecs
 import numpy as np
-from hypothesis import given, settings  # noqa: F401
+from hypothesis import assume, given, settings  # noqa: F401
 from hypothesis.strategies import SearchStrategy
 
 import zarr
@@ -344,3 +345,136 @@ def make_request(start: int, length: int) -> RangeByteRequest:
     )
     key_tuple = st.tuples(keys, byte_ranges)
     return st.lists(key_tuple, min_size=1, max_size=10)
+
+
+def simple_text():
+    """A strategy for generating simple text strings."""
+    return st.text(st.characters(min_codepoint=32, max_codepoint=126), min_size=1, max_size=10)
+
+
+def simple_attrs():
+    """A strategy for generating simple attribute dictionaries."""
+    return st.dictionaries(
+        simple_text(),
+        st.one_of(
+            st.integers(),
+            st.floats(allow_nan=False, allow_infinity=False),
+            st.booleans(),
+            simple_text(),
+        ),
+    )
+
+
+def array_shapes(min_dims=1, max_dims=3, max_len=100):
+    """A strategy for generating array shapes."""
+    return st.lists(
+        st.integers(min_value=1, max_value=max_len), min_size=min_dims, max_size=max_dims
+    )
+
+
+# def zarr_compressors():
+#     """A strategy for generating Zarr compressors."""
+#     return st.sampled_from([None, Blosc(), GZip(), Zstd(), LZ4()])
+
+
+# def zarr_codecs():
+#     """A strategy for generating Zarr codecs."""
+#     return st.sampled_from([BytesCodec(), Blosc(), GZip(), Zstd(), LZ4()])
+
+
+def zarr_filters():
+    """A strategy for generating Zarr filters."""
+    return st.lists(
+        st.just(numcodecs.Delta(dtype="i4")), min_size=0, max_size=2
+    )  # Example filter, expand as needed
+
+
+def zarr_storage_transformers():
+    """A strategy for generating Zarr storage transformers."""
+    return st.lists(
+        st.dictionaries(
+            simple_text(), st.one_of(st.integers(), st.floats(), st.booleans(), simple_text())
+        ),
+        min_size=0,
+        max_size=2,
+    )
+
+
+@st.composite
+def array_metadata_v2(draw: st.DrawFn) -> ArrayV2Metadata:
+    """Generates valid ArrayV2Metadata objects for property-based testing."""
+    dims = draw(st.integers(min_value=1, max_value=3))  # Limit dimensions for complexity
+    shape = tuple(draw(array_shapes(min_dims=dims, max_dims=dims, max_len=100)))
+    max_chunk_len = max(shape) if shape else 100
+    chunks = tuple(
+        draw(
+            st.lists(
+                st.integers(min_value=1, max_value=max_chunk_len), min_size=dims, max_size=dims
+            )
+        )
+    )
+
+    # Validate shape and chunks relationship
+    assume(all(c <= s for s, c in zip(shape, chunks, strict=False)))  # Chunk size must be <= shape
+
+    dtype = draw(v2_dtypes())
+    fill_value = draw(st.one_of([st.none(), npst.from_dtype(dtype)]))
+    order = draw(st.sampled_from(["C", "F"]))
+    dimension_separator = draw(st.sampled_from([".", "/"]))
+    # compressor = draw(zarr_compressors())
+    filters = tuple(draw(zarr_filters())) if draw(st.booleans()) else None
+    attributes = draw(simple_attrs())
+
+    # Construct the metadata object.  Type hints are crucial here for correctness.
+    return ArrayV2Metadata(
+        shape=shape,
+        dtype=dtype,
+        chunks=chunks,
+        fill_value=fill_value,
+        order=order,
+        dimension_separator=dimension_separator,
+        #    compressor=compressor,
+        filters=filters,
+        attributes=attributes,
+    )
+
+
+@st.composite
+def array_metadata_v3(draw: st.DrawFn) -> ArrayV3Metadata:
+    """Generates valid ArrayV3Metadata objects for property-based testing."""
+    dims = draw(st.integers(min_value=1, max_value=3))
+    shape = tuple(draw(array_shapes(min_dims=dims, max_dims=dims, max_len=100)))
+    max_chunk_len = max(shape) if shape else 100
+    chunks = tuple(
+        draw(
+            st.lists(
+                st.integers(min_value=1, max_value=max_chunk_len), min_size=dims, max_size=dims
+            )
+        )
+    )
+    assume(all(c <= s for s, c in zip(shape, chunks, strict=False)))
+
+    dtype = draw(v3_dtypes())
+    fill_value = draw(npst.from_dtype(dtype))
+    chunk_grid = RegularChunkGrid(chunks)  # Ensure chunks is passed as tuple.
+    chunk_key_encoding = DefaultChunkKeyEncoding(separator="/")  # Or st.sampled_from(["/", "."])
+    # codecs = tuple(draw(st.lists(zarr_codecs(), min_size=0, max_size=3)))
+    attributes = draw(simple_attrs())
+    dimension_names = (
+        tuple(draw(st.lists(st.one_of(st.none(), simple_text()), min_size=dims, max_size=dims)))
+        if draw(st.booleans())
+        else None
+    )
+    storage_transformers = tuple(draw(zarr_storage_transformers()))
+
+    return ArrayV3Metadata(
+        shape=shape,
+        data_type=dtype,
+        chunk_grid=chunk_grid,
+        chunk_key_encoding=chunk_key_encoding,
+        fill_value=fill_value,
+        #    codecs=codecs,
+        attributes=attributes,
+        dimension_names=dimension_names,
+        storage_transformers=storage_transformers,
+    )

tests/test_properties.py

@@ -1,3 +1,7 @@
+import dataclasses
+import json
+
+import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
 
@@ -10,10 +14,12 @@
 from hypothesis import assume, given
 
 from zarr.abc.store import Store
+from zarr.core.common import ZARRAY_JSON, ZATTRS_JSON
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
 from zarr.core.sync import sync
 from zarr.testing.strategies import (
     array_metadata,
+    array_metadata_v2,
     arrays,
     basic_indices,
     numpy_arrays,
@@ -23,6 +29,60 @@
 )
 
 
+def deep_equal(a, b):
+    """Deep equality check w/ NaN e to handle array metadata serialization and deserialization behaviors"""
+    if isinstance(a, (complex, np.complexfloating)) and isinstance(
+        b, (complex, np.complexfloating)
+    ):
+        # Convert to Python float to force standard NaN handling.
+        a_real, a_imag = float(a.real), float(a.imag)
+        b_real, b_imag = float(b.real), float(b.imag)
+        # If both parts are NaN, consider them equal.
+        if np.isnan(a_real) and np.isnan(b_real):
+            real_eq = True
+        else:
+            real_eq = a_real == b_real
+        if np.isnan(a_imag) and np.isnan(b_imag):
+            imag_eq = True
+        else:
+            imag_eq = a_imag == b_imag
+        return real_eq and imag_eq
+
+    # Handle floats (including numpy floating types) and treat NaNs as equal.
+    if isinstance(a, (float, np.floating)) and isinstance(b, (float, np.floating)):
+        if np.isnan(a) and np.isnan(b):
+            return True
+        return a == b
+
+    # Handle numpy.datetime64 values, treating NaT as equal.
+    if isinstance(a, np.datetime64) and isinstance(b, np.datetime64):
+        if np.isnat(a) and np.isnat(b):
+            return True
+        return a == b
+
+    # Handle numpy arrays.
+    if isinstance(a, np.ndarray) and isinstance(b, np.ndarray):
+        if a.shape != b.shape:
+            return False
+        # Compare elementwise.
+        return all(deep_equal(x, y) for x, y in zip(a.flat, b.flat, strict=False))
+
+    # Handle dictionaries.
+    if isinstance(a, dict) and isinstance(b, dict):
+        if set(a.keys()) != set(b.keys()):
+            return False
+        return all(deep_equal(a[k], b[k]) for k in a)
+
+    # Handle lists and tuples.
+    if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)):
+        if len(a) != len(b):
+            return False
+        return all(deep_equal(x, y) for x, y in zip(a, b, strict=False))
+
+    # Fallback to default equality.
+    return a == b
+
+
 @given(data=st.data(), zarr_format=zarr_formats)
 def test_roundtrip(data: st.DataObject, zarr_format: int) -> None:
     nparray = data.draw(numpy_arrays(zarr_formats=st.just(zarr_format)))
@@ -135,3 +195,44 @@ async def test_roundtrip_array_metadata(
 #     nparray = data.draw(np_arrays)
 #     zarray = data.draw(arrays(arrays=st.just(nparray)))
 #     assert_array_equal(nparray, zarray[:])
+
+
+@given(array_metadata_v2())
+def test_v2meta_roundtrip(metadata):
+    buffer_dict = metadata.to_buffer_dict(prototype=default_buffer_prototype())
+    zarray_dict = json.loads(buffer_dict[ZARRAY_JSON].to_bytes().decode())
+    zattrs_dict = json.loads(buffer_dict[ZATTRS_JSON].to_bytes().decode())
+
+    # zattrs and zarray are separate in v2, we have to add attributes back prior to `from_dict`
+    zarray_dict["attributes"] = zattrs_dict
+
+    metadata_roundtripped = ArrayV2Metadata.from_dict(zarray_dict)
+
+    # Convert both metadata instances to dictionaries.
+    orig = dataclasses.asdict(metadata)
+    rt = dataclasses.asdict(metadata_roundtripped)
+
+    assert deep_equal(orig, rt), f"Roundtrip mismatch:\nOriginal: {orig}\nRoundtripped: {rt}"
+
+
+@given(npst.from_dtype(dtype=np.dtype("float64"), allow_nan=True, allow_infinity=True))
+def test_v2meta_nan_and_infinity(fill_value):
+    metadata = ArrayV2Metadata(
+        shape=[10],
+        dtype=np.dtype("float64"),
+        chunks=[5],
+        fill_value=fill_value,
+        order="C",
+    )
+
+    buffer_dict = metadata.to_buffer_dict(prototype=default_buffer_prototype())
+    zarray_dict = json.loads(buffer_dict[ZARRAY_JSON].to_bytes().decode())
+
+    if np.isnan(fill_value):
+        assert zarray_dict["fill_value"] == "NaN"
+    elif np.isinf(fill_value) and fill_value > 0:
+        assert zarray_dict["fill_value"] == "Infinity"
+    elif np.isinf(fill_value):
+        assert zarray_dict["fill_value"] == "-Infinity"
+    else:
+        assert zarray_dict["fill_value"] == fill_value