Merge pull request #77 from d-v-b/fill-value-fix

- Ensure that the JSON form of an array's fill value is used for the `fill_value` fields for v2 and v3 `ArraySpec` classes.
- Add the parametric equality method `like` to v3 models
- Refactor `to_zarr` for v2 and v3 models to use identical logic. Note that this is a breaking change, as `to_zarr` for an `ArraySpec` longer takes `**kwargs`. In its place, the `config` keyword argument has been added, that specifies the runtime configuration of the array.
- Ensure that `str` is allowed in the union of possible `fill_value` types. This supports string dtypes
- Run tests against zarr-python 3.0.10 and 3.1.0

Author: d-v-b

.github/workflows/test.yml

@@ -38,6 +38,7 @@ jobs:
         pip install hatch
     - name: Run Tests
       run: |
+        hatch run test.py${{ matrix.python-version }}-${{ matrix.zarr-version }}:list-env
         hatch run test.py${{ matrix.python-version }}-${{ matrix.zarr-version }}:test-cov
     - name: Upload coverage
       uses: codecov/codecov-action@v5

docs/index.md

@@ -39,7 +39,7 @@ print(spec.model_dump())
             'order': 'C',
             'filters': None,
             'dimension_separator': '.',
-            'compressor': {'id': 'zstd', 'level': 0, 'checksum': False},
+            'compressor': {'id': 'zstd', 'level': 0},
         }
     },
 }

pyproject.toml

@@ -44,14 +44,18 @@ build.hooks.vcs.version-file = "src/pydantic_zarr/_version.py"
 
 [tool.hatch.envs.test]
 features = ["test"]
+dependencies = [
+    "zarr~={matrix:zarr}",
+]
 
 [tool.hatch.envs.test.scripts]
 test = "pytest tests/test_pydantic_zarr/"
 test-cov = "pytest --cov-config=pyproject.toml --cov=pkg --cov-report html --cov=src tests/test_pydantic_zarr"
+list-env = "pip list"
 
 [[tool.hatch.envs.test.matrix]]
 python = ["3.11", "3.12", "3.13"]
-zarr_python = ["3.0.10", "3.1.0"]
+zarr = ["3.0.10", "3.1.0"]
 
 [tool.hatch.envs.docs]
 features = ['docs']

src/pydantic_zarr/core.py

@@ -1,27 +1,38 @@
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
     Literal,
     TypeAlias,
+    overload,
 )
 
 import numpy as np
 import numpy.typing as npt
-import zarr
 from pydantic import BaseModel, ConfigDict
+from zarr.core.sync import sync
 from zarr.core.sync_group import get_node
+from zarr.storage._common import make_store_path
 
 if TYPE_CHECKING:
-    from zarr.abc.store import Store
+    import zarr
+    from zarr.storage._common import StoreLike
 
 IncEx: TypeAlias = set[int] | set[str] | dict[int, Any] | dict[str, Any] | None
 
 AccessMode: TypeAlias = Literal["w", "w+", "r", "a"]
 
 
+@overload
+def tuplify_json(obj: Mapping) -> Mapping: ...
+
+
+@overload
+def tuplify_json(obj: list) -> tuple: ...
+
+
 def tuplify_json(obj: object) -> object:
     """
     Recursively converts lists within a Python object to tuples.
@@ -38,21 +49,38 @@ class StrictBase(BaseModel):
     model_config = ConfigDict(frozen=True, extra="forbid")
 
 
-def stringify_dtype(value: npt.DTypeLike) -> str:
+def parse_dtype_v2(value: npt.DTypeLike) -> str | list[tuple[Any, ...]]:
     """
-    Convert a `numpy.dtype` object into a `str`.
+    Convert the input to a NumPy dtype and either return the ``str`` attribute of that
+    object or, if the dtype is a structured dtype, return the fields of that dtype as a list
+    of tuples.
 
     Parameters
     ----------
-    value : `npt.DTypeLike`
-        Some object that can be coerced to a numpy dtype
+    value : npt.DTypeLike
+        A value that can be converted to a NumPy dtype.
 
     Returns
     -------
 
-    A numpy dtype string representation of `value`.
+    A Zarr V2-compatible encoding of the dtype.
+
+    References
+    ----------
+    See the [Zarr V2 specification](https://zarr-specs.readthedocs.io/en/latest/v2/v2.0.html#data-type-encoding)
+    for more details on this encoding of data types.
     """
-    return np.dtype(value).str
+    # Assume that a non-string sequence represents a the Zarr V2 JSON form of a structured dtype.
+    if isinstance(value, Sequence) and not isinstance(value, str):
+        return [tuple(v) for v in value]
+    else:
+        np_dtype = np.dtype(value)
+        if np_dtype.fields is not None:
+            # This is a structured dtype, which must be converted to a list of tuples. Note that
+            # this function recurses, because a structured dtype is parametrized by other dtypes.
+            return [(k, parse_dtype_v2(v[0])) for k, v in np_dtype.fields.items()]
+        else:
+            return np_dtype.str
 
 
 def ensure_member_name(data: Any) -> str:
@@ -92,15 +120,16 @@ def model_like(a: BaseModel, b: BaseModel, exclude: IncEx = None, include: IncEx
 
 # TODO: expose contains_array and contains_group as public functions in zarr-python
 # and replace these custom implementations
-def contains_array(store: Store, path: str) -> bool:
-    try:
-        return isinstance(get_node(store, path, zarr_format=2), zarr.Array)
-    except FileNotFoundError:
-        return False
-
-
-def contains_group(store: Store, path: str) -> bool:
+def maybe_node(
+    store: StoreLike, path: str, *, zarr_format: Literal[2, 3]
+) -> zarr.Array | zarr.Group | None:
+    """
+    Return the array or group found at the store / path, if an array or group exists there.
+    Otherwise return None.
+    """
+    # convert the storelike store argument to a Zarr store
+    spath = sync(make_store_path(store, path=path))
     try:
-        return isinstance(get_node(store, path, zarr_format=2), zarr.Group)
+        return get_node(spath.store, spath.path, zarr_format=zarr_format)
     except FileNotFoundError:
-        return False
+        return None

src/pydantic_zarr/v2.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import json
 import math
 from collections.abc import Mapping
+from importlib.metadata import version
 from typing import (
     TYPE_CHECKING,
     Annotated,
@@ -17,28 +19,31 @@
     overload,
 )
 
-import numcodecs
 import numpy as np
 import numpy.typing as npt
 import zarr
 from numcodecs.abc import Codec
+from packaging.version import Version
 from pydantic import AfterValidator, BaseModel, field_validator, model_validator
 from pydantic.functional_validators import BeforeValidator
+from zarr.core.array import Array, AsyncArray
 from zarr.core.metadata import ArrayV2Metadata
+from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
+from zarr.storage._common import make_store_path
 
 from pydantic_zarr.core import (
     IncEx,
     StrictBase,
-    contains_array,
-    contains_group,
     ensure_key_no_path,
+    maybe_node,
     model_like,
-    stringify_dtype,
+    parse_dtype_v2,
 )
 
 if TYPE_CHECKING:
     from zarr.abc.store import Store
+    from zarr.core.array_spec import ArrayConfigParams
 
 TBaseAttr: TypeAlias = Mapping[str, object] | BaseModel
 TBaseItem: TypeAlias = Union["GroupSpec", "ArraySpec"]
@@ -49,7 +54,18 @@
 TAttr = TypeVar("TAttr", bound=TBaseAttr)
 TItem = TypeVar("TItem", bound=TBaseItem)
 
-DtypeStr = Annotated[str, BeforeValidator(stringify_dtype)]
+DtypeStr = Annotated[str, BeforeValidator(parse_dtype_v2)]
+
+BoolFillValue = bool
+IntFillValue = int
+# todo: introduce a type that represents hexadecimal representations of floats
+FloatFillValue = Literal["Infinity", "-Infinity", "NaN"] | float
+ComplexFillValue = tuple[FloatFillValue, FloatFillValue]
+RawFillValue = tuple[int, ...]
+
+FillValue = (
+    BoolFillValue | IntFillValue | FloatFillValue | ComplexFillValue | RawFillValue | str | None
+)
 
 DimensionSeparator = Literal[".", "/"]
 MemoryOrder = Literal["C", "F"]
@@ -155,8 +171,8 @@ class ArraySpec(NodeSpec, Generic[TAttr]):
     attributes: TAttr = cast(TAttr, {})
     shape: tuple[int, ...]
     chunks: tuple[int, ...]
-    dtype: DtypeStr
-    fill_value: int | float | None = 0
+    dtype: DtypeStr | list[tuple[Any, ...]]
+    fill_value: FillValue = 0
     order: MemoryOrder = "C"
     filters: list[CodecDict] | None = None
     dimension_separator: Annotated[
@@ -285,7 +301,7 @@ def from_array(
 
         return cls(
             shape=shape_actual,
-            dtype=stringify_dtype(dtype_actual),
+            dtype=parse_dtype_v2(dtype_actual),
             chunks=chunks_actual,
             attributes=attributes_actual,
             fill_value=fill_value_actual,
@@ -322,40 +338,25 @@ def from_zarr(cls, array: zarr.Array) -> Self:
             msg = "Array is not a Zarr format 2 array"
             raise TypeError(msg)
 
-        if len(array.compressors):
-            compressor = array.compressors[0]
-            if TYPE_CHECKING:
-                # TODO: overload array.compressors in zarr-python and remove this type check
-                assert isinstance(compressor, Codec)
-            compressor_dict = compressor.get_config()
+        if Version(version("zarr")) < Version("3.1.0"):
+            from zarr.core.buffer import default_buffer_prototype
+
+            stored_meta = array.metadata.to_buffer_dict(prototype=default_buffer_prototype())
+            meta_json = json.loads(stored_meta[".zarray"].to_bytes()) | {
+                "attributes": array.attrs.asdict()
+            }
         else:
-            compressor_dict = None
+            meta_json = array.metadata.to_dict()
 
-        return cls(
-            shape=array.shape,
-            chunks=array.chunks,
-            dtype=str(array.dtype),
-            # explicitly cast to numpy type and back to python
-            # so that int 0 isn't serialized as 0.0
-            fill_value=(
-                array.dtype.type(array.fill_value).tolist()
-                if array.fill_value is not None
-                else array.fill_value
-            ),
-            order=array.order,
-            filters=array.filters,
-            dimension_separator=array.metadata.dimension_separator,
-            compressor=compressor_dict,
-            attributes=array.attrs.asdict(),
-        )
+        return cls.model_validate(meta_json)
 
     def to_zarr(
         self,
         store: Store,
         path: str,
         *,
         overwrite: bool = False,
-        **kwargs: Any,
+        config: ArrayConfigParams | None = None,
     ) -> zarr.Array:
         """
         Serialize an `ArraySpec` to a Zarr array at a specific path in a Zarr store. This operation
@@ -369,36 +370,32 @@ def to_zarr(
             The location of the array inside the store.
         overwrite : bool, default = False
             Whether to overwrite existing objects in storage to create the Zarr array.
-        **kwargs : Any
-            Additional keyword arguments are passed to `zarr.create`.
+        config : ArrayConfigParams | None, default = None
+            An instance of `ArrayConfigParams` that defines the runtime configuration for the array.
 
         Returns
         -------
         zarr.Array
             A Zarr array that is structurally identical to `self`.
         """
-        spec_dict = self.model_dump()
-        attrs = spec_dict.pop("attributes")
-        if self.compressor is not None:
-            spec_dict["compressor"] = numcodecs.get_codec(spec_dict["compressor"])
-        if self.filters is not None:
-            spec_dict["filters"] = [numcodecs.get_codec(f) for f in spec_dict["filters"]]
-        if contains_array(store, path):
-            extant_array = zarr.open_array(store, path=path, zarr_format=2)
-
-            if not self.like(extant_array):
-                if not overwrite:
-                    raise ContainsArrayError(store, path)
+        store_path = sync(make_store_path(store, path=path))
+
+        extant_node = maybe_node(store, path, zarr_format=2)
+        if isinstance(extant_node, zarr.Array):
+            if not self.like(extant_node) and not overwrite:
+                raise ContainsArrayError(store, path)
             else:
+                # If there's an existing array that is identical to the model, and overwrite is False,
+                # we can just return that existing array.
                 if not overwrite:
-                    # extant_array is read-only, so we make a new array handle that
-                    # takes **kwargs
-                    return zarr.open_array(
-                        store=extant_array.store, path=extant_array.path, zarr_format=2, **kwargs
-                    )
-        result = zarr.create(store=store, path=path, overwrite=overwrite, **spec_dict, **kwargs)
-        result.attrs.put(attrs)
-        return result
+                    return extant_node
+        if isinstance(extant_node, zarr.Group) and not overwrite:
+            raise ContainsGroupError(store, path)
+
+        meta: ArrayV2Metadata = ArrayV2Metadata.from_dict(self.model_dump())
+        async_array = AsyncArray(metadata=meta, store_path=store_path, config=config)
+        sync(async_array._save_metadata(meta))
+        return Array(_async_array=async_array)
 
     def like(
         self,
@@ -568,28 +565,34 @@ def to_zarr(
         """
         spec_dict = self.model_dump(exclude={"members": True})
         attrs = spec_dict.pop("attributes")
-        if contains_group(store, path):
-            extant_group = zarr.group(store, path=path, zarr_format=2)
-            if not self.like(extant_group):
+        extant_node = maybe_node(store, path, zarr_format=2)
+        if isinstance(extant_node, zarr.Group):
+            if not self.like(extant_node):
                 if not overwrite:
+                    """
                     msg = (
                         f"A group already exists at path {path}. "
                         "That group is structurally dissimilar to the group you are trying to store."
                         "Call to_zarr with overwrite=True to overwrite that group."
                     )
-                    raise ContainsGroupError(msg)
+                    """
+                    # Zarr's contains group error uses questionable design and doesn't take a message
+                    raise ContainsGroupError(store, path)
             else:
                 if not overwrite:
                     # if the extant group is structurally identical to self, and overwrite is false,
                     # then just return the extant group
-                    return extant_group
+                    return extant_node
 
-        elif contains_array(store, path) and not overwrite:
+        elif isinstance(extant_node, zarr.Array) and not overwrite:
+            """
             msg = (
                 f"An array already exists at path {path}. "
                 "Call to_zarr with overwrite=True to overwrite the array."
             )
-            raise ContainsArrayError(msg)
+            """
+            # Zarr's contains array error uses questionable design and doesn't take a message
+            raise ContainsArrayError(store, path)
         else:
             zarr.create_group(store=store, overwrite=overwrite, path=path, zarr_format=2)