feat(attributes): warn when writing attributes that are not valid JSON

changes/4400.feature.md

@@ -0,0 +1 @@
+Writing attributes with non-string keys, which JSON stores as strings, now emits a `ZarrFutureWarning`: this will become an error in a future version. Two new config options control how attributes that are not valid JSON are written: `attributes.non_string_keys` (default `"warn"`) and `attributes.non_finite_floats` for `NaN` and infinite values (default `"allow"`). Each takes `"allow"`, `"warn"` or `"raise"`.

docs/user-guide/attributes.md

@@ -54,6 +54,28 @@ print(sorted(z.attrs))
 Internally Zarr uses JSON to store array and group attributes, so attribute
 values must be JSON serializable.
 
+Two kinds of value that Python's `json` module accepts are not valid JSON, and
+the config controls how Zarr writes them. Each option takes `"allow"` (write
+without comment), `"warn"` (write, and emit a `ZarrFutureWarning`) or `"raise"`
+(refuse to write):
+
+- Non-string keys, such as `{1: "a"}`, are written as strings, so they read
+  back as different keys. `attributes.non_string_keys` defaults to `"warn"`,
+  and this will become an error in a future version of Zarr.
+- `NaN` and infinite floats are written as non-standard `NaN` / `Infinity`
+  literals that other JSON parsers may reject.
+  `attributes.non_finite_floats` defaults to `"allow"`. Set it to `"raise"` to
+  make sure the documents Zarr writes are valid JSON.
+
+```python exec="true" session="attributes" source="above" result="ansi"
+arr = zarr.create_array(store="memory://attributes-nan-demo", shape=(1,), dtype="f8")
+with zarr.config.set({"attributes.non_finite_floats": "raise"}):
+    try:
+        arr.attrs["scale"] = float("nan")
+    except ValueError as e:
+        print(e)
+```
+
 When copying a Zarr array with [`zarr.from_array`][], its attributes are
 deep-copied by default so nested dictionaries and lists are independent of the
 source. Deeply nested attributes can raise `RecursionError` during this copy,

docs/user-guide/config.md

@@ -40,6 +40,7 @@ Configuration options include the following:
 - Whether empty chunks are written to storage `array.write_empty_chunks`
 - Enable experimental rectilinear chunks `array.rectilinear_chunks`
 - Whether missing chunks are filled with the array's fill value on read `array.read_missing_chunks` (default `True`). Set to `False` to raise a [`ChunkNotFoundError`][zarr.errors.ChunkNotFoundError] instead.
+- How attributes that are not valid JSON are written: `attributes.non_string_keys` (default `"warn"`) and `attributes.non_finite_floats` (default `"allow"`). Each takes `"allow"`, `"warn"` or `"raise"`. See [Working with attributes](attributes.md).
 - Async and threading options, e.g. `async.concurrency` and `threading.max_workers`
 - Selections of implementations of codecs, codec pipelines and buffers
 - Enabling GPU support with `zarr.config.enable_gpu()`. See [GPU support](gpu.md) for more.

src/zarr/core/config.py

@@ -107,6 +107,13 @@ def enable_gpu(self) -> ConfigSet:
             "async": {"concurrency": 10, "timeout": None},
             "threading": {"max_workers": None},
             "json_indent": 2,
+            "attributes": {
+                # How attributes that are not valid JSON are handled when they are
+                # written: "allow", "warn", or "raise". See
+                # zarr.core.metadata.common.check_attributes_json.
+                "non_string_keys": "warn",
+                "non_finite_floats": "allow",
+            },
             "codec_pipeline": {
                 # FusedCodecPipeline is the faster synchronous pipeline, but it stays
                 # opt-in for now so behavior is unchanged for existing users. Early

src/zarr/core/group.py

@@ -48,6 +48,7 @@
 from zarr.core.dtype import parse_data_type
 from zarr.core.json_parse import parse_field
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
+from zarr.core.metadata.common import check_attributes_json
 from zarr.core.metadata.io import save_metadata
 from zarr.core.sync import SyncMixin, sync
 from zarr.errors import (
@@ -359,8 +360,13 @@ class GroupMetadata(Metadata):
     def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         indent = config.get("json_indent")
         if self.zarr_format == 3:
-            return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)}
+            # `to_dict` first: it rejects non-string top-level keys with a TypeError,
+            # which must keep taking precedence over the attributes policy.
+            zarr_json = self.to_dict()
+            check_attributes_json(self.attributes)
+            return {ZARR_JSON: json_to_buffer(zarr_json, prototype=prototype, indent=indent)}
         else:
+            check_attributes_json(self.attributes)
             items = {
                 ZGROUP_JSON: json_to_buffer(
                     {"zarr_format": self.zarr_format}, prototype=prototype, indent=indent

src/zarr/core/metadata/common.py

@@ -1,13 +1,130 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import math
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Final, Literal, cast
+
+from zarr.core.config import config
+from zarr.errors import ZarrFutureWarning
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
     from zarr.core.common import JSON
 
+AttributeJSONPolicy = Literal["allow", "warn", "raise"]
+ATTRIBUTE_JSON_POLICIES: Final = ("allow", "warn", "raise")
+
+# Key types that `json.dumps` silently converts to strings. Other non-string
+# key types make `json.dumps` raise, so they need no check here.
+_COERCED_KEY_TYPES: Final = (int, float, bool, type(None))
+
+# The root of the zarr package, used to attribute warnings to the caller outside zarr.
+_ZARR_PACKAGE_ROOT: Final = str(Path(__file__).parents[2])
+
 
 def parse_attributes(data: dict[str, JSON] | None) -> dict[str, JSON]:
     if data is None:
         return {}
 
     return dict(data)
+
+
+def _iter_non_json(
+    value: Any, path: tuple[object, ...] = ()
+) -> Iterator[tuple[Literal["non_string_key", "non_finite_float"], tuple[object, ...]]]:
+    """Yield `(kind, path)` for every part of `value` that `json.dumps` accepts
+    but that is not round-trippable JSON.
+
+    `non_string_key` marks a dict key that `json.dumps` converts to a string;
+    `non_finite_float` marks a NaN or infinite float, which `json.dumps` writes
+    as a non-standard `NaN` / `Infinity` literal.
+    """
+    if isinstance(value, dict):
+        for key, item in value.items():
+            if isinstance(key, _COERCED_KEY_TYPES):
+                yield "non_string_key", (*path, key)
+            yield from _iter_non_json(item, (*path, key))
+    elif isinstance(value, list | tuple):
+        for index, item in enumerate(value):
+            yield from _iter_non_json(item, (*path, index))
+    elif isinstance(value, float) and not math.isfinite(value):
+        yield "non_finite_float", path
+
+
+def _format_path(path: tuple[object, ...]) -> str:
+    return "attributes" + "".join(f"[{part!r}]" for part in path)
+
+
+def _get_policy(key: str) -> AttributeJSONPolicy:
+    policy = config.get(key)
+    if policy not in ATTRIBUTE_JSON_POLICIES:
+        raise ValueError(
+            f"Invalid value for config option {key!r}: {policy!r}. "
+            f"Expected one of {ATTRIBUTE_JSON_POLICIES}."
+        )
+    return cast("AttributeJSONPolicy", policy)
+
+
+def check_attributes_json(attributes: dict[str, JSON]) -> None:
+    """Apply the configured policy to attributes that are not valid JSON.
+
+    Called when attributes are serialized for writing. Two kinds of value are
+    checked, each under its own config option:
+
+    - `attributes.non_string_keys`: dict keys that are not strings. JSON object
+      keys are strings, so these are written as strings and read back as
+      different keys. Defaults to `"warn"`; this will become an error in a
+      future version of zarr.
+    - `attributes.non_finite_floats`: NaN and infinite floats. JSON cannot
+      represent them, so they are written as non-standard `NaN` / `Infinity`
+      literals that other JSON parsers may reject. Defaults to `"allow"`.
+
+    Each option is `"allow"` (write without comment), `"warn"` (write, and emit
+    a `ZarrFutureWarning`) or `"raise"` (refuse to write).
+    """
+    problems: dict[str, list[tuple[object, ...]]] = {
+        "non_string_key": [],
+        "non_finite_float": [],
+    }
+    for kind, path in _iter_non_json(attributes):
+        problems[kind].append(path)
+
+    if problems["non_string_key"]:
+        policy = _get_policy("attributes.non_string_keys")
+        if policy != "allow":
+            paths = ", ".join(_format_path(p) for p in problems["non_string_key"])
+            detail = (
+                f"Attribute keys must be strings, but these keys are not: {paths}. "
+                "JSON object keys are strings, so these keys are written to the store as "
+                "strings and read back as different keys."
+            )
+            if policy == "raise":
+                raise TypeError(detail)
+            warnings.warn(
+                f"{detail} This will be an error in a future version of zarr. Convert the "
+                "keys to strings, or set the config option 'attributes.non_string_keys' to "
+                "'raise' to opt in to the error now, or to 'allow' to keep the current "
+                "behavior without this warning.",
+                ZarrFutureWarning,
+                skip_file_prefixes=(_ZARR_PACKAGE_ROOT,),
+            )
+
+    if problems["non_finite_float"]:
+        policy = _get_policy("attributes.non_finite_floats")
+        if policy != "allow":
+            paths = ", ".join(_format_path(p) for p in problems["non_finite_float"])
+            detail = (
+                f"Attribute values must be valid JSON, but these values are NaN or infinite: "
+                f"{paths}. JSON cannot represent them, so they are written as non-standard "
+                "NaN / Infinity literals that other JSON parsers may reject."
+            )
+            if policy == "raise":
+                raise ValueError(detail)
+            warnings.warn(
+                f"{detail} Set the config option 'attributes.non_finite_floats' to 'raise' "
+                "to refuse to write them, or to 'allow' to write them without this warning.",
+                ZarrFutureWarning,
+                skip_file_prefixes=(_ZARR_PACKAGE_ROOT,),
+            )

src/zarr/core/metadata/v2.py

@@ -42,7 +42,7 @@
 )
 from zarr.core.config import config, parse_indexing_order
 from zarr.core.json_parse import parse_field
-from zarr.core.metadata.common import parse_attributes
+from zarr.core.metadata.common import check_attributes_json, parse_attributes
 
 
 class ArrayV2MetadataDict(TypedDict):
@@ -142,6 +142,7 @@ def shards(self) -> tuple[int, ...] | None:
     def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         zarray_dict = self.to_dict()
         zattrs_dict = zarray_dict.pop("attributes", {})
+        check_attributes_json(self.attributes)
         indent = config.get("json_indent")
         return {
             ZARRAY_JSON: json_to_buffer(zarray_dict, prototype=prototype, indent=indent),

src/zarr/core/metadata/v3.py

@@ -35,7 +35,7 @@
 from zarr.core.dtype import VariableLengthUTF8, ZDType, get_data_type_from_json
 from zarr.core.dtype.common import check_dtype_spec_v3
 from zarr.core.json_parse import parse_field
-from zarr.core.metadata.common import parse_attributes
+from zarr.core.metadata.common import check_attributes_json, parse_attributes
 from zarr.errors import MetadataValidationError, NodeTypeValidationError
 from zarr.registry import get_codec_class
 
@@ -627,7 +627,9 @@ def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str:
 
     def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         indent = config.get("json_indent")
-        return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)}
+        zarr_json = self.to_dict()
+        check_attributes_json(self.attributes)
+        return {ZARR_JSON: json_to_buffer(zarr_json, prototype=prototype, indent=indent)}
 
     @classmethod
     def from_dict(cls, data: dict[str, JSON]) -> Self: