Support unlistable stores for all v04 groups (#278)

* Support unlistable stores for Well

* Move get_array_paths onto attribute objects

* Move _to_zarr to shared file

* Add different API for optional array paths

* Implement checking group paths

* Remove old print statements

* Fix HCS with unlistable stores

* Add comment

* Use _from_zarr with Labels

* Fix Well.from_zarr() docstring

Author: dstansby

src/ome_zarr_models/_v06/image.py

@@ -54,9 +54,8 @@ def from_zarr(cls, group: zarr.Group) -> Self:  # type: ignore[override]
         members_tree_flat: dict[str, AnyGroupSpec | AnyArraySpec] = {}
         for multiscale in multi_meta.multiscales:
             for dataset in multiscale.datasets:
-                array_path = f"{group.path}/{dataset.path}"
                 array_spec = check_array_path(
-                    group, array_path, expected_zarr_version=3
+                    group, dataset.path, expected_zarr_version=3
                 )
                 members_tree_flat["/" + dataset.path] = array_spec
 

src/ome_zarr_models/base.py

@@ -2,6 +2,7 @@
 from typing import Literal
 
 from pydantic import BaseModel, ConfigDict
+from pydantic_zarr.v2 import AnyGroupSpec
 
 
 class BaseAttrs(BaseModel):
@@ -21,6 +22,38 @@ class BaseAttrs(BaseModel):
         frozen=True,
     )
 
+    def get_array_paths(self) -> list[str]:
+        """
+        Get a list of all array paths expected and required to live in the Group
+        with these attributes.
+        """
+        return []
+
+    def get_optional_array_paths(self) -> list[str]:
+        """
+        Get a list of all array paths expected but not required to live in the Group
+        with these attributes.
+        """
+        return []
+
+    def get_group_paths(self) -> dict[str, type[AnyGroupSpec]]:
+        """
+        Get a list of all group paths expected and required to live in the Group
+        with these attributes.
+
+        Must return a dictionary mapping paths to their GroupSpec class.
+        """
+        return {}
+
+    def get_optional_group_paths(self) -> dict[str, type[AnyGroupSpec]]:
+        """
+        Get a list of all group paths expected but not required to live in the Group
+        with these attributes.
+
+        Must return a dictionary mapping paths to their GroupSpec class.
+        """
+        return {}
+
 
 class BaseGroup(ABC):
     """

src/ome_zarr_models/common/validation.py

@@ -88,31 +88,94 @@ def check_array_path(
         If the array doesn't exist, or the array is not the expected Zarr version.
     """
     try:
-        array = zarr.open_array(store=group.store, path=array_path, mode="r")
-        array_spec: AnyArraySpecv2 | AnyArraySpecv3
-        if array.metadata.zarr_format == 2:
-            if expected_zarr_version == 3:
-                raise ValueError("Expected Zarr v3 array, but got v2 array")
-            array_spec = ArraySpecv2.from_zarr(array)
-        else:
-            if expected_zarr_version == 2:
-                raise ValueError("Expected Zarr v2 array, but got v3 array")
-            array_spec = ArraySpecv3.from_zarr(array)
+        array = zarr.open_array(store=group.store_path, path=array_path, mode="r")
     except FileNotFoundError as e:
         msg = (
             f"Expected to find an array at {array_path}, but no array was found there."
         )
         raise ValueError(msg) from e
-    except zarr.errors.ContainsGroupError as e:
+    except (zarr.errors.ContainsGroupError, zarr.errors.NodeTypeValidationError) as e:
         msg = (
             f"Expected to find an array at {array_path}, "
             "but a group was found there instead."
         )
         raise ValueError(msg) from e
 
+    array_spec: AnyArraySpecv2 | AnyArraySpecv3
+    if array.metadata.zarr_format == 2:
+        if expected_zarr_version == 3:
+            raise ValueError("Expected Zarr v3 array, but got v2 array")
+        array_spec = ArraySpecv2.from_zarr(array)
+    else:
+        if expected_zarr_version == 2:
+            raise ValueError("Expected Zarr v2 array, but got v3 array")
+        array_spec = ArraySpecv3.from_zarr(array)
+
     return array_spec
 
 
+@overload
+def check_group_path(
+    group: zarr.Group,
+    group_path: str,
+    *,
+    expected_zarr_version: Literal[2],
+) -> AnyGroupSpecv2: ...
+
+
+@overload
+def check_group_path(
+    group: zarr.Group,
+    group_path: str,
+    *,
+    expected_zarr_version: Literal[3],
+) -> AnyGroupSpecv3: ...
+
+
+def check_group_path(
+    group: zarr.Group,
+    group_path: str,
+    *,
+    expected_zarr_version: Literal[2, 3],
+) -> AnyGroupSpecv2 | AnyGroupSpecv3:
+    """
+    Check if a group exists at a given path in a group.
+
+    Returns
+    -------
+    GroupSpec
+        If the path exists, it's GroupSpec is returned.
+
+    Raises
+    ------
+    ValueError
+        If the group doesn't exist, or the group is not the expected Zarr version.
+    """
+    try:
+        group = zarr.open_group(store=group.store_path, path=group_path, mode="r")
+    except FileNotFoundError as e:
+        msg = f"Expected to find a group at {group_path}, but no group was found there."
+        raise FileNotFoundError(msg) from e
+    except zarr.errors.ContainsArrayError as e:
+        msg = (
+            f"Expected to find an group at {group_path}, "
+            "but an array was found there instead."
+        )
+        raise zarr.errors.ContainsArrayError(msg) from e
+
+    group_spec: AnyGroupSpecv2 | AnyGroupSpecv3
+    if group.metadata.zarr_format == 2:
+        if expected_zarr_version == 3:
+            raise ValueError("Expected Zarr v3 array, but got v2 array")
+        group_spec = GroupSpecv2.from_zarr(group, depth=0)
+    else:
+        if expected_zarr_version == 2:
+            raise ValueError("Expected Zarr v2 array, but got v3 array")
+        group_spec = GroupSpecv3.from_zarr(group, depth=0)
+
+    return group_spec
+
+
 def check_length(
     sequence: Sequence[T], *, valid_lengths: Sequence[int], variable_name: str
 ) -> None:

src/ome_zarr_models/common/well_types.py

@@ -32,7 +32,9 @@ class WellMeta(BaseAttrs):
     Metadata for a single well.
     """
 
-    images: Annotated[list[WellImage], AfterValidator(unique_items_validator)]
+    images: Annotated[list[WellImage], AfterValidator(unique_items_validator)] = Field(
+        ..., description="Images within a well"
+    )
     version: str | None = Field(None, description="Version of the well specification")
 
     def get_acquisition_paths(self) -> dict[int, list[str]]:

src/ome_zarr_models/v04/_shared.py

@@ -0,0 +1,80 @@
+from typing import TYPE_CHECKING, Any, TypeVar
+
+import zarr
+from pydantic_zarr.v2 import AnyGroupSpec, GroupSpec
+
+from ome_zarr_models.base import BaseAttrs
+from ome_zarr_models.common.validation import (
+    check_array_path,
+    check_group_path,
+)
+from ome_zarr_models.v04.base import BaseGroupv04
+
+TCls = TypeVar("TCls", bound=BaseGroupv04[Any])
+TAttrs = TypeVar("TAttrs", bound=BaseAttrs)
+
+if TYPE_CHECKING:
+    from pydantic_zarr.v2 import AnyArraySpec, AnyGroupSpec
+
+
+def _from_zarr(
+    group: zarr.Group,
+    group_cls: type[TCls],
+    attrs_cls: type[TAttrs],
+) -> TCls:
+    """
+    Create a GroupSpec from a potentially unlistable Zarr group.
+
+    This uses methods on the attribute class to get required and optional
+    paths to ararys and groups, and then manually constructs the GroupSpec
+    from those paths.
+
+    Parameters
+    ----------
+    group :
+        Zarr group to create GroupSpec from.
+    group_cls :
+        Class of the Group to return.
+    attrs_cls :
+        Attributes class.
+    """
+    # on unlistable storage backends, the members of this group will be {}
+    group_spec: AnyGroupSpec = GroupSpec.from_zarr(group, depth=0)
+    attributes = attrs_cls.model_validate(group_spec.attributes)
+
+    members_tree_flat: dict[str, AnyGroupSpec | AnyArraySpec] = {}
+
+    # Required array paths
+    for array_path in attrs_cls.get_array_paths(attributes):
+        array_spec = check_array_path(group, array_path, expected_zarr_version=2)
+        members_tree_flat["/" + array_path] = array_spec
+
+    # Optional array paths
+    for array_path in attrs_cls.get_optional_array_paths(attributes):
+        try:
+            array_spec = check_array_path(group, array_path, expected_zarr_version=2)
+        except ValueError:
+            continue
+        members_tree_flat["/" + array_path] = array_spec
+
+    # Required group paths
+    required_groups = attrs_cls.get_group_paths(attributes)
+    for group_path in required_groups:
+        group_spec = check_group_path(group, group_path, expected_zarr_version=2)
+        group_flat = required_groups[group_path].from_zarr(group[group_path]).to_flat()  # type: ignore[arg-type]
+        for path in group_flat:
+            members_tree_flat["/" + group_path + path] = group_flat[path]
+
+    # Optional group paths
+    optional_groups = attrs_cls.get_optional_group_paths(attributes)
+    for group_path in optional_groups:
+        try:
+            group_spec = check_group_path(group, group_path, expected_zarr_version=2)
+        except FileNotFoundError:
+            continue
+        group_flat = optional_groups[group_path].from_zarr(group[group_path]).to_flat()  # type: ignore[arg-type]
+        for path in group_flat:
+            members_tree_flat["/" + group_path + path] = group_flat[path]
+
+    members_normalized: AnyGroupSpec = GroupSpec.from_flat(members_tree_flat)
+    return group_cls(members=members_normalized.members, attributes=attributes)

src/ome_zarr_models/v04/hcs.py

@@ -1,11 +1,13 @@
 from collections.abc import Generator, Mapping
 from typing import Self
 
+import zarr
 from pydantic import model_validator
-from pydantic_zarr.v2 import GroupSpec
+from pydantic_zarr.v2 import AnyGroupSpec, GroupSpec
 
 from ome_zarr_models.base import BaseAttrs
 from ome_zarr_models.common.well import WellGroupNotFoundError
+from ome_zarr_models.v04._shared import _from_zarr
 from ome_zarr_models.v04.base import BaseGroupv04
 from ome_zarr_models.v04.plate import Plate
 from ome_zarr_models.v04.well import Well
@@ -20,12 +22,37 @@ class HCSAttrs(BaseAttrs):
 
     plate: Plate
 
+    def get_optional_group_paths(self) -> dict[str, type[AnyGroupSpec]]:  # noqa: D102
+        return {well.path: Well for well in self.plate.wells}
+
 
 class HCS(BaseGroupv04[HCSAttrs]):
     """
     An OME-Zarr high-content screening (HCS) dataset representing a single plate.
     """
 
+    @classmethod
+    def from_zarr(cls, group: zarr.Group) -> Self:  # type: ignore[override]
+        """
+        Create an OME-Zarr image model from a `zarr.Group`.
+
+        Parameters
+        ----------
+        group : zarr.Group
+            A Zarr group that has valid OME-Zarr image metadata.
+        """
+        hcs = _from_zarr(group, cls, HCSAttrs)
+        # Traverse all the Well groups, which themselves contain Image groups
+        hcs_flat = hcs.to_flat()
+        for well in hcs.attributes.plate.wells:
+            well_group = group[well.path]
+            well_group_flat = Well.from_zarr(well_group).to_flat()  # type: ignore[arg-type]
+            for path in well_group_flat:
+                hcs_flat["/" + well.path + path] = well_group_flat[path]
+
+        hcs_unflat: AnyGroupSpec = GroupSpec.from_flat(hcs_flat)
+        return cls(attributes=hcs_unflat.attributes, members=hcs_unflat.members)
+
     @model_validator(mode="after")
     def _check_valid_acquisitions(self) -> Self:
         """

src/ome_zarr_models/v04/image.py

@@ -2,14 +2,12 @@
 
 from typing import TYPE_CHECKING, Self
 
-import zarr
-import zarr.errors
 from pydantic import Field, JsonValue, model_validator
 from pydantic_zarr.v2 import AnyArraySpec, AnyGroupSpec, GroupSpec
 
 from ome_zarr_models.base import BaseAttrs
 from ome_zarr_models.common.coordinate_transformations import _build_transforms
-from ome_zarr_models.common.validation import check_array_path
+from ome_zarr_models.v04._shared import _from_zarr
 from ome_zarr_models.v04.axes import Axis
 from ome_zarr_models.v04.base import BaseGroupv04
 from ome_zarr_models.v04.labels import Labels
@@ -19,6 +17,8 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    import zarr
+
 
 __all__ = ["Image", "ImageAttrs"]
 
@@ -35,6 +35,16 @@ class ImageAttrs(BaseAttrs):
     )
     omero: Omero | None = None
 
+    def get_array_paths(self) -> list[str]:  # noqa: D102
+        paths = []
+        for multiscale in self.multiscales:
+            for dataset in multiscale.datasets:
+                paths.append(dataset.path)
+        return paths
+
+    def get_optional_group_paths(self) -> dict[str, type[AnyGroupSpec]]:  # noqa: D102
+        return {"labels": Labels}
+
 
 class Image(BaseGroupv04[ImageAttrs]):
     """
@@ -51,31 +61,7 @@ def from_zarr(cls, group: zarr.Group) -> Self:  # type: ignore[override]
         group : zarr.Group
             A Zarr group that has valid OME-Zarr image metadata.
         """
-        # on unlistable storage backends, the members of this group will be {}
-        group_spec: AnyGroupSpec = GroupSpec.from_zarr(group, depth=0)
-
-        multi_meta = ImageAttrs.model_validate(group_spec.attributes)
-        members_tree_flat: dict[str, AnyGroupSpec | AnyArraySpec] = {}
-        for multiscale in multi_meta.multiscales:
-            for dataset in multiscale.datasets:
-                array_path = f"{group.path}/{dataset.path}"
-                array_spec = check_array_path(
-                    group, array_path, expected_zarr_version=2
-                )
-                members_tree_flat["/" + dataset.path] = array_spec
-
-        try:
-            labels_group = zarr.open_group(store=group.store, path="labels", mode="r")
-            members_tree_flat["/labels"] = GroupSpec.from_zarr(labels_group, depth=0)
-        except zarr.errors.GroupNotFoundError:
-            pass
-
-        members_normalized: AnyGroupSpec = GroupSpec.from_flat(members_tree_flat)
-
-        group_spec = group_spec.model_copy(
-            update={"members": members_normalized.members}
-        )
-        return cls(**group_spec.model_dump())
+        return _from_zarr(group, cls, ImageAttrs)
 
     @classmethod
     def new(