convert inf, -inf, nan to JSON (#3280)

* convert numpy inf, -inf, nan to JSON

* changelog

Author: d-v-b

changes/3280.fix.rst

@@ -0,0 +1,2 @@
+Fix a regression introduced in 3.1.0 that prevented ``inf``, ``-inf``, and ``nan`` values
+from being stored in ``attributes``.

src/zarr/core/group.py

@@ -336,7 +336,7 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         if self.zarr_format == 3:
             return {
                 ZARR_JSON: prototype.buffer.from_bytes(
-                    json.dumps(self.to_dict(), indent=json_indent, allow_nan=False).encode()
+                    json.dumps(self.to_dict(), indent=json_indent, allow_nan=True).encode()
                 )
             }
         else:
@@ -345,7 +345,7 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
                     json.dumps({"zarr_format": self.zarr_format}, indent=json_indent).encode()
                 ),
                 ZATTRS_JSON: prototype.buffer.from_bytes(
-                    json.dumps(self.attributes, indent=json_indent, allow_nan=False).encode()
+                    json.dumps(self.attributes, indent=json_indent, allow_nan=True).encode()
                 ),
             }
             if self.consolidated_metadata:
@@ -373,7 +373,7 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
 
                 items[ZMETADATA_V2_JSON] = prototype.buffer.from_bytes(
                     json.dumps(
-                        {"metadata": d, "zarr_consolidated_format": 1}, allow_nan=False
+                        {"metadata": d, "zarr_consolidated_format": 1}, allow_nan=True
                     ).encode()
                 )
 

src/zarr/core/metadata/v2.py

@@ -132,10 +132,10 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         json_indent = config.get("json_indent")
         return {
             ZARRAY_JSON: prototype.buffer.from_bytes(
-                json.dumps(zarray_dict, indent=json_indent, allow_nan=False).encode()
+                json.dumps(zarray_dict, indent=json_indent, allow_nan=True).encode()
             ),
             ZATTRS_JSON: prototype.buffer.from_bytes(
-                json.dumps(zattrs_dict, indent=json_indent, allow_nan=False).encode()
+                json.dumps(zattrs_dict, indent=json_indent, allow_nan=True).encode()
             ),
         }
 

src/zarr/core/metadata/v3.py

@@ -288,7 +288,7 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         d = self.to_dict()
         return {
             ZARR_JSON: prototype.buffer.from_bytes(
-                json.dumps(d, allow_nan=False, indent=json_indent).encode()
+                json.dumps(d, allow_nan=True, indent=json_indent).encode()
             )
         }
 

tests/conftest.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import math
 import os
 import pathlib
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
@@ -442,3 +444,21 @@ def skip_object_dtype(dtype: ZDType[Any, Any]) -> None:
             "type resolution"
         )
         pytest.skip(msg)
+
+
+def nan_equal(a: object, b: object) -> bool:
+    """
+    Convenience function for equality comparison between two values ``a`` and ``b``, that might both
+    be NaN. Returns True if both ``a`` and ``b`` are NaN, otherwise returns a == b
+    """
+    if math.isnan(a) and math.isnan(b):  # type: ignore[arg-type]
+        return True
+    return a == b
+
+
+def deep_nan_equal(a: object, b: object) -> bool:
+    if isinstance(a, Mapping) and isinstance(b, Mapping):
+        return all(deep_nan_equal(a[k], b[k]) for k in a)
+    if isinstance(a, Sequence) and isinstance(b, Sequence):
+        return all(deep_nan_equal(a[i], b[i]) for i in range(len(a)))
+    return nan_equal(a, b)

tests/test_attributes.py

@@ -1,18 +1,26 @@
+import json
+from typing import Any
+
+import numpy as np
 import pytest
 
 import zarr.core
 import zarr.core.attributes
 import zarr.storage
+from tests.conftest import deep_nan_equal
+from zarr.core.common import ZarrFormat
 
 
-def test_put() -> None:
+@pytest.mark.parametrize("zarr_format", [2, 3])
+@pytest.mark.parametrize(
+    "data", [{"inf": np.inf, "-inf": -np.inf, "nan": np.nan}, {"a": 3, "c": 4}]
+)
+def test_put(data: dict[str, Any], zarr_format: ZarrFormat) -> None:
     store = zarr.storage.MemoryStore()
-    attrs = zarr.core.attributes.Attributes(
-        zarr.Group.from_store(store, attributes={"a": 1, "b": 2})
-    )
-    attrs.put({"a": 3, "c": 4})
-    expected = {"a": 3, "c": 4}
-    assert dict(attrs) == expected
+    attrs = zarr.core.attributes.Attributes(zarr.Group.from_store(store, zarr_format=zarr_format))
+    attrs.put(data)
+    expected = json.loads(json.dumps(data, allow_nan=True))
+    assert deep_nan_equal(dict(attrs), expected)
 
 
 def test_asdict() -> None:

tests/test_dtype/test_npy/test_common.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 import base64
-import math
 import re
 import sys
 from typing import TYPE_CHECKING, Any, get_args
 
 import numpy as np
 import pytest
 
+from tests.conftest import nan_equal
 from zarr.core.dtype.common import ENDIANNESS_STR, JSONFloatV2, SpecialFloatStrings
 from zarr.core.dtype.npy.common import (
     NumpyEndiannessStr,
@@ -35,16 +35,6 @@
     from zarr.core.common import JSON, ZarrFormat
 
 
-def nan_equal(a: object, b: object) -> bool:
-    """
-    Convenience function for equality comparison between two values ``a`` and ``b``, that might both
-    be NaN. Returns True if both ``a`` and ``b`` are NaN, otherwise returns a == b
-    """
-    if math.isnan(a) and math.isnan(b):  # type: ignore[arg-type]
-        return True
-    return a == b
-
-
 json_float_v2_roundtrip_cases: tuple[tuple[JSONFloatV2, float | np.floating[Any]], ...] = (
     ("Infinity", float("inf")),
     ("Infinity", np.inf),