fix array metadata dicts and refactor to_dict test

Author: d-v-b

tests/test_metadata/test_v3.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 import json
 import re
 from typing import TYPE_CHECKING, Literal
@@ -10,6 +11,7 @@
 from zarr.codecs.bytes import BytesCodec
 from zarr.core.buffer import default_buffer_prototype
 from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding, V2ChunkKeyEncoding
+from zarr.core.common import ArrayMetadataJSON_V3, NamedConfig
 from zarr.core.config import config
 from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.core.dtype.npy.string import _NUMPY_SUPPORTS_VLEN_STRING
@@ -110,83 +112,84 @@ def test_parse_fill_value_invalid_type_sequence(fill_value: Any, dtype_str: str)
         dtype_instance.from_json_scalar(fill_value, zarr_format=3)
 
 
-@pytest.mark.parametrize("chunk_grid", ["regular"])
-@pytest.mark.parametrize("attributes", [None, {"foo": "bar"}])
-@pytest.mark.parametrize("codecs", [[BytesCodec(endian=None)]])
+@pytest.mark.parametrize("chunk_grid", [{"name": "regular", "configuration": {"chunk_shape": (1, 1, 1)}}])
+@pytest.mark.parametrize("codecs", [({"name" : "bytes"},)])
 @pytest.mark.parametrize("fill_value", [0, 1])
-@pytest.mark.parametrize("chunk_key_encoding", ["v2", "default"])
-@pytest.mark.parametrize("dimension_separator", [".", "/", None])
-@pytest.mark.parametrize("dimension_names", ["nones", "strings", "missing"])
-@pytest.mark.parametrize("storage_transformers", [None, ()])
+@pytest.mark.parametrize("data_type", ["int8", "uint8"])
+@pytest.mark.parametrize("chunk_key_encoding", [
+    {"name": "v2", "configuration": {"separator": "."}},
+    {"name": "v2", "configuration": {"separator": "/"}},
+    {"name": "v2"},
+    {"name": "default", "configuration": {"separator": "."}},
+    {"name": "default", "configuration": {"separator": "/"}},
+    {"name": "default"},
+    ])
+@pytest.mark.parametrize("attributes", ["unset", {"foo": "bar"}])
+@pytest.mark.parametrize("dimension_names", [(None, None, None), ('a','b', None), "unset"])
+@pytest.mark.parametrize("storage_transformers", [(), "unset"])
 def test_metadata_to_dict(
-    chunk_grid: str,
+    chunk_grid: NamedConfig[str, Mapping[str, object]],
     codecs: list[Codec],
+    data_type: str,
     fill_value: Any,
-    chunk_key_encoding: Literal["v2", "default"],
-    dimension_separator: Literal[".", "/"] | None,
-    dimension_names: Literal["nones", "strings", "missing"],
-    attributes: dict[str, Any] | None,
-    storage_transformers: tuple[dict[str, JSON]] | None,
+    chunk_key_encoding: NamedConfig[str, Mapping[str, object]],
+    dimension_names: tuple[str | None, ...] | Literal["unset"],
+    attributes: dict[str, Any] | Literal['unset'],
+    storage_transformers: tuple[dict[str, JSON]] | Literal["unset"],
 ) -> None:
     shape = (1, 2, 3)
-    data_type_str = "uint8"
-    if chunk_grid == "regular":
-        cgrid = {"name": "regular", "configuration": {"chunk_shape": (1, 1, 1)}}
-
-    cke: dict[str, Any]
-    cke_name_dict = {"name": chunk_key_encoding}
-    if dimension_separator is not None:
-        cke = cke_name_dict | {"configuration": {"separator": dimension_separator}}
-    else:
-        cke = cke_name_dict
-    dnames: tuple[str | None, ...] | None
-
-    if dimension_names == "strings":
-        dnames = tuple(map(str, range(len(shape))))
-    elif dimension_names == "missing":
-        dnames = None
-    elif dimension_names == "nones":
-        dnames = (None,) * len(shape)
-
-    metadata_dict = {
-        "zarr_format": 3,
-        "node_type": "array",
-        "shape": shape,
-        "chunk_grid": cgrid,
-        "data_type": data_type_str,
-        "chunk_key_encoding": cke,
-        "codecs": tuple(c.to_dict() for c in codecs),
-        "fill_value": fill_value,
-        "storage_transformers": storage_transformers,
-    }
 
-    if attributes is not None:
-        metadata_dict["attributes"] = attributes
-    if dnames is not None:
-        metadata_dict["dimension_names"] = dnames
+    # These are the fields in the array metadata document that are optional
+    not_required = {}
 
-    metadata = ArrayV3Metadata.from_dict(metadata_dict)
-    observed = metadata.to_dict()
-    expected = metadata_dict.copy()
+    if dimension_names != "unset":
+        not_required["dimension_names"] = dimension_names
 
-    # if unset or None or (), storage_transformers gets normalized to ()
-    assert observed["storage_transformers"] == ()
-    observed.pop("storage_transformers")
-    expected.pop("storage_transformers")
+    if storage_transformers != "unset":
+        not_required["storage_transformers"] = storage_transformers
 
-    if attributes is None:
-        assert observed["attributes"] == {}
-        observed.pop("attributes")
+    if attributes != "unset":
+        not_required["attributes"] = attributes
 
-    if dimension_separator is None:
-        if chunk_key_encoding == "default":
-            expected_cke_dict = DefaultChunkKeyEncoding(separator="/").to_dict()
+    source_dict = {
+        "zarr_format": 3,
+        "node_type": "array",
+        "shape": shape,
+        "chunk_grid": chunk_grid,
+        "data_type": data_type,
+        "chunk_key_encoding": chunk_key_encoding,
+        "codecs": codecs,
+        "fill_value": fill_value,
+    } | not_required
+
+    metadata = ArrayV3Metadata.from_dict(source_dict)
+    parsed_dict = metadata.to_dict()
+
+    for k,v in parsed_dict.items():
+        if k in source_dict:
+            if k == 'chunk_key_encoding':
+                assert v['name'] == chunk_key_encoding['name']
+                if chunk_key_encoding['name'] == 'v2':
+                    if "configuration" in chunk_key_encoding:
+                        if "separator" in chunk_key_encoding['configuration']:
+                            assert v['configuration']['separator'] == chunk_key_encoding['configuration']['separator']
+                    else:
+                        assert v["configuration"]["separator"] == "."
+                elif chunk_key_encoding['name'] == 'default':
+                    if "configuration" in chunk_key_encoding:
+                        if "separator" in chunk_key_encoding['configuration']:
+                            assert v['configuration']['separator'] == chunk_key_encoding['configuration']['separator']
+                    else:
+                        assert v["configuration"]["separator"] == "/"
+            else:
+                assert source_dict[k] == v
         else:
-            expected_cke_dict = V2ChunkKeyEncoding(separator=".").to_dict()
-        assert observed["chunk_key_encoding"] == expected_cke_dict
-        observed.pop("chunk_key_encoding")
-        expected.pop("chunk_key_encoding")
-    assert observed == expected
+            if k == 'attributes':
+                assert v == {}
+            elif k == 'storage_transformers':
+                assert v == ()
+            else:
+                assert v is None
 
 
 @pytest.mark.parametrize("indent", [2, 4, None])
@@ -201,14 +204,14 @@ def test_json_indent(indent: int):
 @pytest.mark.parametrize("precision", ["ns", "D"])
 async def test_datetime_metadata(fill_value: int, precision: str) -> None:
     dtype = DateTime64(unit=precision)
-    metadata_dict = {
+    metadata_dict: ArrayMetadataJSON_V3 = {
         "zarr_format": 3,
         "node_type": "array",
         "shape": (1,),
         "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": (1,)}},
         "data_type": dtype.to_json(zarr_format=3),
-        "chunk_key_encoding": {"name": "default", "separator": "."},
-        "codecs": (BytesCodec(),),
+        "chunk_key_encoding": {"name": "default", "configuration": {"separator": "."}},
+        "codecs": ({"name": "bytes"},),
         "fill_value": dtype.to_json_scalar(
             dtype.to_native_dtype().type(fill_value, dtype.unit), zarr_format=3
         ),
@@ -225,13 +228,13 @@ async def test_datetime_metadata(fill_value: int, precision: str) -> None:
     ("data_type", "fill_value"), [("uint8", {}), ("int32", [0, 1]), ("float32", "foo")]
 )
 async def test_invalid_fill_value_raises(data_type: str, fill_value: float) -> None:
-    metadata_dict = {
+    metadata_dict: ArrayMetadataJSON_V3 = {
         "zarr_format": 3,
         "node_type": "array",
         "shape": (1,),
         "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": (1,)}},
         "data_type": data_type,
-        "chunk_key_encoding": {"name": "default", "separator": "."},
+        "chunk_key_encoding": {"name": "default", "configuration": {"separator": "."}},
         "codecs": ({"name": "bytes"},),
         "fill_value": fill_value,  # this is not a valid fill value for uint8
     }
@@ -242,13 +245,13 @@ async def test_invalid_fill_value_raises(data_type: str, fill_value: float) -> N
 
 @pytest.mark.parametrize("fill_value", [("NaN"), "Infinity", "-Infinity"])
 async def test_special_float_fill_values(fill_value: str) -> None:
-    metadata_dict = {
+    metadata_dict: ArrayMetadataJSON_V3 = {
         "zarr_format": 3,
         "node_type": "array",
         "shape": (1,),
         "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": (1,)}},
         "data_type": "float64",
-        "chunk_key_encoding": {"name": "default", "separator": "."},
+        "chunk_key_encoding": {"name": "default", "configuration": {"separator": "."}},
         "codecs": [{"name": "bytes"}],
         "fill_value": fill_value,  # this is not a valid fill value for uint8
     }