Remove redundancies, fix integral handling

Author: moradology

src/zarr/core/metadata/v2.py

@@ -2,7 +2,7 @@
 
 import base64
 import warnings
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Sequence
 from enum import Enum
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, TypedDict, cast
@@ -109,29 +109,6 @@ def shards(self) -> ChunkCoords | None:
         return None
 
     def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
-        def _serialize_fill_value(fv: Any) -> JSON:
-            if self.fill_value is None:
-                pass
-            elif self.dtype.kind in "SV":
-                # There's a relationship between self.dtype and self.fill_value
-                # that mypy isn't aware of. The fact that we have S or V dtype here
-                # means we should have a bytes-type fill_value.
-                fv = base64.standard_b64encode(cast(bytes, self.fill_value)).decode("ascii")
-            elif isinstance(fv, np.datetime64):
-                if np.isnat(fv):
-                    fv = "NaT"
-                else:
-                    fv = np.datetime_as_string(fv)
-            elif isinstance(fv, numbers.Real):
-                float_fv = float(fv)
-                if np.isnan(float_fv):
-                    fv = "NaN"
-                elif np.isinf(float_fv):
-                    fv = "Infinity" if float_fv > 0 else "-Infinity"
-            elif isinstance(fv, numbers.Complex):
-                fv = [_serialize_fill_value(fv.real), _serialize_fill_value(fv.imag)]
-            return cast(JSON, fv)
-
         def _json_convert(
             o: Any,
         ) -> Any:
@@ -170,7 +147,7 @@ def _json_convert(
             raise TypeError
 
         zarray_dict = self.to_dict()
-        zarray_dict["fill_value"] = _serialize_fill_value(zarray_dict["fill_value"])
+        zarray_dict["fill_value"] = _serialize_fill_value(self.fill_value, self.dtype)
         zattrs_dict = zarray_dict.pop("attributes", {})
         json_indent = config.get("json_indent")
         return {
@@ -185,7 +162,7 @@ def _json_convert(
         }
 
     @classmethod
-    def from_dict(cls, data: dict[str, JSON]) -> ArrayV2Metadata:
+    def from_dict(cls, data: dict[str, Any]) -> ArrayV2Metadata:
         # Make a copy to protect the original from modification.
         _data = data.copy()
         # Check that the zarr_format attribute is correct.
@@ -213,7 +190,7 @@ def from_dict(cls, data: dict[str, JSON]) -> ArrayV2Metadata:
 
         _data = {k: v for k, v in _data.items() if k in expected}
 
-        return cls(**cast(Mapping[str, Any], _data))
+        return cls(**_data)
 
     def to_dict(self) -> dict[str, JSON]:
         zarray_dict = super().to_dict()
@@ -315,7 +292,7 @@ def parse_metadata(data: ArrayV2Metadata) -> ArrayV2Metadata:
     return data
 
 
-def parse_structured_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
+def _parse_structured_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
     """Handle structured dtype/fill value pairs"""
     try:
         if isinstance(fill_value, list):
@@ -354,7 +331,10 @@ def parse_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
     if fill_value is None or dtype.hasobject:
         pass
     elif dtype.fields is not None:
-        fill_value = parse_structured_fill_value(fill_value, dtype)
+        # the dtype is structured (has multiple fields), so the fill_value might be a
+        # compound value (e.g., a tuple or dict) that needs field-wise processing.
+        # We use parse_structured_fill_value to correctly convert each component.
+        fill_value = _parse_structured_fill_value(fill_value, dtype)
     elif not isinstance(fill_value, np.void) and fill_value == 0:
         # this should be compatible across numpy versions for any array type, including
         # structured arrays
@@ -369,16 +349,9 @@ def parse_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
             )
     elif dtype.kind in "SV" and isinstance(fill_value, str):
         fill_value = base64.standard_b64decode(fill_value)
-    elif np.issubdtype(dtype, np.datetime64):
-        if fill_value == "NaT":
-            fill_value = np.array("NaT", dtype=dtype)[()]
-        else:
-            fill_value = np.array(fill_value, dtype=dtype)[()]
     elif dtype.kind == "c" and isinstance(fill_value, list) and len(fill_value) == 2:
         complex_val = complex(float(fill_value[0]), float(fill_value[1]))
         fill_value = np.array(complex_val, dtype=dtype)[()]
-    elif dtype.kind in "f" and fill_value in {"NaN", "Infinity", "-Infinity"}:
-        fill_value = np.array(fill_value, dtype=dtype)[()]
     else:
         try:
             if isinstance(fill_value, bytes) and dtype.kind == "V":
@@ -394,6 +367,39 @@ def parse_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
     return fill_value
 
 
+def _serialize_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> JSON:
+    serialized: JSON
+
+    if fill_value is None:
+        serialized = None
+    elif dtype.kind in "SV":
+        # There's a relationship between dtype and fill_value
+        # that mypy isn't aware of. The fact that we have S or V dtype here
+        # means we should have a bytes-type fill_value.
+        serialized = base64.standard_b64encode(cast(bytes, fill_value)).decode("ascii")
+    elif isinstance(fill_value, np.datetime64):
+        serialized = np.datetime_as_string(fill_value)
+    elif isinstance(fill_value, numbers.Integral):
+        serialized = int(fill_value)
+    elif isinstance(fill_value, numbers.Real):
+        float_fv = float(fill_value)
+        if np.isnan(float_fv):
+            serialized = "NaN"
+        elif np.isinf(float_fv):
+            serialized = "Infinity" if float_fv > 0 else "-Infinity"
+        else:
+            serialized = float_fv
+    elif isinstance(fill_value, numbers.Complex):
+        serialized = [
+            _serialize_fill_value(fill_value.real, dtype),
+            _serialize_fill_value(fill_value.imag, dtype),
+        ]
+    else:
+        serialized = fill_value
+
+    return serialized
+
+
 def _default_fill_value(dtype: np.dtype[Any]) -> Any:
     """
     Get the default fill value for a type.

tests/test_properties.py

@@ -1,6 +1,7 @@
 import dataclasses
 import json
 import numbers
+from typing import Any
 
 import numpy as np
 import pytest
@@ -12,7 +13,7 @@
 
 import hypothesis.extra.numpy as npst
 import hypothesis.strategies as st
-from hypothesis import assume, given
+from hypothesis import assume, given, settings
 
 from zarr.abc.store import Store
 from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON
@@ -30,7 +31,7 @@
 )
 
 
-def deep_equal(a, b):
+def deep_equal(a: Any, b: Any) -> bool:
     """Deep equality check with handling of special cases for array metadata classes"""
     if isinstance(a, (complex, np.complexfloating)) and isinstance(
         b, (complex, np.complexfloating)
@@ -100,6 +101,8 @@ def test_array_creates_implicit_groups(array):
             )
 
 
+# bump deadline from 200 to 300 to avoid (rare) intermittent timeouts
+@settings(deadline=300)
 @given(data=st.data())
 def test_basic_indexing(data: st.DataObject) -> None:
     zarray = data.draw(simple_arrays())
@@ -236,7 +239,7 @@ def test_roundtrip_array_metadata_from_json(data: st.DataObject, zarr_format: in
 #     assert_array_equal(nparray, zarray[:])
 
 
-def serialized_float_is_valid(serialized):
+def serialized_float_is_valid(serialized: numbers.Real | str) -> bool:
     """
     Validate that the serialized representation of a float conforms to the spec.
 

tests/test_v2.py

@@ -15,7 +15,7 @@
 from zarr import config
 from zarr.abc.store import Store
 from zarr.core.buffer.core import default_buffer_prototype
-from zarr.core.metadata.v2 import parse_structured_fill_value
+from zarr.core.metadata.v2 import _parse_structured_fill_value
 from zarr.core.sync import sync
 from zarr.storage import MemoryStore, StorePath
 
@@ -366,7 +366,7 @@ def test_structured_dtype_roundtrip(fill_value, tmp_path) -> None:
 def test_parse_structured_fill_value_valid(
     fill_value: Any, dtype: np.dtype[Any], expected_result: Any
 ) -> None:
-    result = parse_structured_fill_value(fill_value, dtype)
+    result = _parse_structured_fill_value(fill_value, dtype)
     assert result.dtype == expected_result.dtype
     assert result == expected_result
     if isinstance(expected_result, np.void):
@@ -396,7 +396,7 @@ def test_parse_structured_fill_value_valid(
 )
 def test_parse_structured_fill_value_invalid(fill_value: Any, dtype: np.dtype[Any]) -> None:
     with pytest.raises(ValueError):
-        parse_structured_fill_value(fill_value, dtype)
+        _parse_structured_fill_value(fill_value, dtype)
 
 
 @pytest.mark.parametrize("fill_value", [None, b"x"], ids=["no_fill", "fill"])