Merge pull request #16 from janelia-cellmap/groupspec_from_zarr_fix

groupspec from zarr fix

Author: d-v-b

src/pydantic_zarr/core.py

@@ -8,6 +8,7 @@
     Optional,
     TypeVar,
     Union,
+    overload,
 )
 from pydantic import BaseModel, root_validator, validator
 from pydantic.generics import GenericModel
@@ -141,7 +142,7 @@ def from_zarr(cls, zarray: zarr.Array):
             filters=zarray.filters,
             dimension_separator=zarray._dimension_separator,
             compressor=zarray.compressor,
-            attrs=dict(zarray.attrs),
+            attrs=zarray.attrs.asdict(),
         )
 
     def to_zarr(
@@ -209,7 +210,7 @@ def from_zarr(cls, group: zarr.Group) -> "GroupSpec[TAttr, TItem]":
             if isinstance(member, zarr.Array):
                 _item = ArraySpec.from_zarr(member)
             elif isinstance(member, zarr.Group):
-                _item = cls.from_zarr(member)
+                _item = GroupSpec.from_zarr(member)
             else:
                 msg = f"""
                 Unparseable object encountered: {type(member)}. Expected zarr.Array or
@@ -218,7 +219,7 @@ def from_zarr(cls, group: zarr.Group) -> "GroupSpec[TAttr, TItem]":
                 raise ValueError(msg)
             members[name] = _item
 
-        result = cls(attrs=dict(group.attrs), members=members)
+        result = cls(attrs=group.attrs.asdict(), members=members)
         return result
 
     def to_zarr(self, store: BaseStore, path: str, overwrite: bool = False):
@@ -259,6 +260,16 @@ def to_zarr(self, store: BaseStore, path: str, overwrite: bool = False):
         return result
 
 
+@overload
+def from_zarr(element: zarr.Array) -> ArraySpec:
+    ...
+
+
+@overload
+def from_zarr(element: zarr.Group) -> GroupSpec:
+    ...
+
+
 def from_zarr(element: Union[zarr.Array, zarr.Group]) -> Union[ArraySpec, GroupSpec]:
     """
     Recursively parse a Zarr group or Zarr array into an ArraySpec or GroupSpec.
@@ -290,7 +301,7 @@ def from_zarr(element: Union[zarr.Array, zarr.Group]) -> Union[ArraySpec, GroupS
                 raise ValueError(msg)
             members[name] = _item
 
-        result = GroupSpec(attrs=dict(element.attrs), members=members)
+        result = GroupSpec(attrs=element.attrs.asdict(), members=members)
         return result
     else:
         msg = f"""

tests/test_core.py

@@ -2,7 +2,7 @@
 import pytest
 import zarr
 from zarr.errors import ContainsGroupError
-from typing import Any, Literal, TypedDict
+from typing import Any, Literal, TypedDict, Union
 import numcodecs
 from pydantic_zarr.core import ArraySpec, GroupSpec, to_zarr, from_zarr
 import numpy as np
@@ -150,12 +150,14 @@ class SubGroupAttrs(TypedDict):
         a: str
         b: float
 
+    SubGroup = GroupSpec[SubGroupAttrs, Any]
+
     class ArrayAttrs(TypedDict):
         scale: list[float]
 
     store = zarr.MemoryStore()
 
-    spec = GroupSpec(
+    spec = GroupSpec[RootAttrs, Union[ArraySpec, SubGroup]](
         attrs=RootAttrs(foo=10, bar=[0, 1, 2]),
         members={
             "s0": ArraySpec(
@@ -178,7 +180,7 @@ class ArrayAttrs(TypedDict):
                 dimension_separator=dimension_separator,
                 attrs=ArrayAttrs(scale=[2.0]),
             ),
-            "subgroup": GroupSpec(attrs=SubGroupAttrs(a="foo", b=1.0)),
+            "subgroup": SubGroup(attrs=SubGroupAttrs(a="foo", b=1.0)),
         },
     )
     # materialize a zarr group, based on the spec
@@ -195,6 +197,11 @@ class ArrayAttrs(TypedDict):
     group2 = to_zarr(spec, store, "/group_a", overwrite=True)
     assert group2 == group
 
+    # again with class methods
+    group3 = spec.to_zarr(store, "/group_b")
+    observed = spec.from_zarr(group3)
+    assert observed == spec
+
 
 def test_shape_chunks():
     for a, b in zip(range(1, 5), range(2, 6)):