Support unlistable stores in v05 and v06 (#284)

* Move _to_zarr helper utility

* Create BaseAttrsv2

* Rename typevar

* Create BaseAttrsv3

* Add path methods to BaseAttrsv3

* Rename typevar

* Add _from_zarr_v3

* Use from_zarrv3 with v05 well

* Import modules instead of classes

* Use from_zarr for v05 image

* Support unlistable stores for v05 HCS

* Support unlistable stores in v06

* Update changelog.

* Remove errant print

* Fix typing errors

Author: dstansby

docs/changelog.md

@@ -4,7 +4,7 @@
 
 ### New Features
 
-- [ome_zarr_models.v04.ImageLabel][] now supports being created from a Zarr array backed by an un-listable store.
+- All models now support being created from Zarr groups in remote HTTP stores, and more generally from any groups stored in any unlistable store.
 
 ## 1.0
 

pyproject.toml

@@ -77,6 +77,7 @@ ignore = [
     "D205",  # 1 blank line required between summary line and description
     "D400",  # First line should end with a period
     "D100",  # Missing docstring in public module
+    "D102",  # Missing docstring in public method
     "D104",  # Missing docstring in public package
     "TC001", # Move application import `...` into a type-checking block. This messes with nested pydantic models
 

src/ome_zarr_models/_utils.py

@@ -2,16 +2,177 @@
 Private utilities.
 """
 
+from __future__ import annotations
+
 from collections import Counter
-from collections.abc import Hashable, Iterable
 from dataclasses import MISSING, fields, is_dataclass
-from typing import TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import pydantic
+import pydantic_zarr.v2
+import pydantic_zarr.v3
 from pydantic import create_model
-from zarr.abc.store import Store
 
-T = TypeVar("T")
+from ome_zarr_models.base import BaseAttrsv2, BaseAttrsv3
+from ome_zarr_models.common.validation import (
+    check_array_path,
+    check_group_path,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Hashable, Iterable
+
+    import zarr
+    from zarr.abc.store import Store
+
+    from ome_zarr_models._v06.base import BaseGroupv06
+    from ome_zarr_models.v04.base import BaseGroupv04
+    from ome_zarr_models.v05.base import BaseGroupv05
+
+TBaseGroupv2 = TypeVar("TBaseGroupv2", bound="BaseGroupv04[Any]")
+TAttrsv2 = TypeVar("TAttrsv2", bound=BaseAttrsv2)
+
+
+def _from_zarr_v2(
+    group: zarr.Group,
+    group_cls: type[TBaseGroupv2],
+    attrs_cls: type[TAttrsv2],
+) -> TBaseGroupv2:
+    """
+    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_in: pydantic_zarr.v2.AnyGroupSpec
+    group_spec_in = pydantic_zarr.v2.GroupSpec.from_zarr(group, depth=0)
+    attributes = attrs_cls.model_validate(group_spec_in.attributes)
+
+    members_tree_flat: dict[
+        str, pydantic_zarr.v2.AnyGroupSpec | pydantic_zarr.v2.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:
+        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:
+            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: pydantic_zarr.v2.AnyGroupSpec = (
+        pydantic_zarr.v2.GroupSpec.from_flat(members_tree_flat)
+    )
+    return group_cls(members=members_normalized.members, attributes=attributes)
+
+
+TBaseGroupv3 = TypeVar("TBaseGroupv3", bound="BaseGroupv05[Any] | BaseGroupv06[Any]")
+TAttrsv3 = TypeVar("TAttrsv3", bound=BaseAttrsv3)
+
+
+def _from_zarr_v3(
+    group: zarr.Group,
+    group_cls: type[TBaseGroupv3],
+    attrs_cls: type[TAttrsv3],
+) -> TBaseGroupv3:
+    """
+    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_in: pydantic_zarr.v3.AnyGroupSpec
+    group_spec_in = pydantic_zarr.v3.GroupSpec.from_zarr(group, depth=0)
+    ome_attributes = attrs_cls.model_validate(group.attrs.asdict()["ome"])
+
+    members_tree_flat: dict[
+        str, pydantic_zarr.v3.AnyGroupSpec | pydantic_zarr.v3.AnyArraySpec
+    ] = {}
+
+    # Required array paths
+    for array_path in attrs_cls.get_array_paths(ome_attributes):
+        array_spec = check_array_path(group, array_path, expected_zarr_version=3)
+        members_tree_flat["/" + array_path] = array_spec
+
+    # Optional array paths
+    for array_path in attrs_cls.get_optional_array_paths(ome_attributes):
+        try:
+            array_spec = check_array_path(group, array_path, expected_zarr_version=3)
+        except ValueError:
+            continue
+        members_tree_flat["/" + array_path] = array_spec
+
+    # Required group paths
+    required_groups = attrs_cls.get_group_paths(ome_attributes)
+    for group_path in required_groups:
+        check_group_path(group, group_path, expected_zarr_version=3)
+        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(ome_attributes)
+    for group_path in optional_groups:
+        try:
+            check_group_path(group, group_path, expected_zarr_version=3)
+        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: pydantic_zarr.v3.AnyGroupSpec = (
+        pydantic_zarr.v3.GroupSpec.from_flat(members_tree_flat)
+    )
+    return group_cls(  # type: ignore[return-value]
+        members=members_normalized.members, attributes=group_spec_in.attributes
+    )
 
 
 def get_store_path(store: Store) -> str:

src/ome_zarr_models/_v06/base.py

@@ -6,15 +6,15 @@
 import pydantic_zarr.v3
 from pydantic import BaseModel
 
-from ome_zarr_models.base import BaseAttrs, BaseGroup
+from ome_zarr_models.base import BaseAttrsv3, BaseGroup
 
 if TYPE_CHECKING:
     import zarr
 
 
-class BaseOMEAttrs(BaseAttrs):
+class BaseOMEAttrs(BaseAttrsv3):
     """
-    Base class for attributes under an OME-Zarr group.
+    Base class for OME-Zarr 0.6 attributes.
     """
 
     version: Literal["0.6"]

src/ome_zarr_models/_v06/hcs.py

@@ -1,16 +1,21 @@
 from collections.abc import Generator, Mapping
-from typing import Self
+from typing import TYPE_CHECKING, Self
 
 # Import needed for pydantic type resolution
 import pydantic_zarr  # noqa: F401
+import zarr
 from pydantic import model_validator
 from pydantic_zarr.v3 import GroupSpec
 
+from ome_zarr_models._utils import _from_zarr_v3
 from ome_zarr_models._v06.base import BaseGroupv06, BaseOMEAttrs
 from ome_zarr_models._v06.plate import Plate
 from ome_zarr_models._v06.well import Well
 from ome_zarr_models.common.well import WellGroupNotFoundError
 
+if TYPE_CHECKING:
+    from pydantic_zarr.v3 import AnyGroupSpec
+
 __all__ = ["HCS", "HCSAttrs"]
 
 
@@ -21,12 +26,38 @@ class HCSAttrs(BaseOMEAttrs):
 
     plate: Plate
 
+    def get_optional_group_paths(self) -> dict[str, type[Well]]:  # type: ignore[override]
+        return {well.path: Well for well in self.plate.wells}
+
 
 class HCS(BaseGroupv06[HCSAttrs]):
     """
     An OME-Zarr high content screening (HCS) dataset.
     """
 
+    @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_v3(group, cls, HCSAttrs)
+        # Traverse all the Well groups, which themselves contain Image groups
+        hcs_flat = hcs.to_flat()
+        for well in hcs.ome_attributes.plate.wells:
+            if well.path in group:
+                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/_v06/image.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Any, Self
+from typing import Self
 
 # Import needed for pydantic type resolution
 import pydantic_zarr  # noqa: F401
@@ -8,12 +8,12 @@
 from pydantic import Field, JsonValue, model_validator
 from pydantic_zarr.v3 import AnyArraySpec, AnyGroupSpec, GroupSpec
 
+from ome_zarr_models._utils import _from_zarr_v3
 from ome_zarr_models._v06.axes import Axis
 from ome_zarr_models._v06.base import BaseGroupv06, BaseOMEAttrs, BaseZarrAttrs
 from ome_zarr_models._v06.labels import Labels
 from ome_zarr_models._v06.multiscales import Dataset, Multiscale
 from ome_zarr_models.common.coordinate_transformations import _build_transforms
-from ome_zarr_models.common.validation import check_array_path
 
 __all__ = ["Image", "ImageAttrs"]
 
@@ -29,6 +29,16 @@ class ImageAttrs(BaseOMEAttrs):
         min_length=1,
     )
 
+    def get_array_paths(self) -> list[str]:
+        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[Labels]]:  # type: ignore[override]
+        return {"labels": Labels}
+
 
 class Image(BaseGroupv06[ImageAttrs]):
     """
@@ -45,33 +55,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: GroupSpec[dict[str, Any], Any] = GroupSpec.from_zarr(group, depth=0)
-
-        if "ome" not in group_spec.attributes:
-            raise RuntimeError(f"Did not find 'ome' key in {group} attributes")
-        multi_meta = ImageAttrs.model_validate(group_spec.attributes["ome"])
-        members_tree_flat: dict[str, AnyGroupSpec | AnyArraySpec] = {}
-        for multiscale in multi_meta.multiscales:
-            for dataset in multiscale.datasets:
-                array_spec = check_array_path(
-                    group, dataset.path, expected_zarr_version=3
-                )
-                members_tree_flat["/" + dataset.path] = array_spec
-
-        try:
-            labels_group = zarr.open_group(store=group.store_path / "labels", mode="r")
-            labels = Labels.from_zarr(labels_group)
-            # members_tree_flat["/labels"] = labels
-            labels_flat = labels.to_flat()
-            for path in labels_flat:
-                members_tree_flat[f"/labels{path}"] = labels_flat[path]
-
-        except zarr.errors.GroupNotFoundError:
-            pass
-
-        members_normalized: AnyGroupSpec = GroupSpec.from_flat(members_tree_flat)
-        return cls(attributes=group_spec.attributes, members=members_normalized.members)
+        return _from_zarr_v3(group, cls, ImageAttrs)
 
     @classmethod
     def new(

src/ome_zarr_models/_v06/well.py

@@ -1,7 +1,11 @@
 # Import needed for pydantic type resolution
-import pydantic_zarr  # noqa: F401
+from typing import Self
 
+import zarr
+
+from ome_zarr_models._utils import _from_zarr_v3
 from ome_zarr_models._v06.base import BaseGroupv06, BaseOMEAttrs
+from ome_zarr_models._v06.image import Image
 from ome_zarr_models._v06.well_types import WellMeta
 
 __all__ = ["Well", "WellAttrs"]
@@ -14,8 +18,23 @@ class WellAttrs(BaseOMEAttrs):
 
     well: WellMeta
 
+    def get_optional_group_paths(self) -> dict[str, type[Image]]:  # type: ignore[override]
+        return {im.path: Image for im in self.well.images}
+
 
 class Well(BaseGroupv06[WellAttrs]):
     """
     An OME-Zarr well dataset.
     """
+
+    @classmethod
+    def from_zarr(cls, group: zarr.Group) -> Self:  # type: ignore[override]
+        """
+        Create an OME-Zarr well model from a `zarr.Group`.
+
+        Parameters
+        ----------
+        group : zarr.Group
+            A Zarr group that has valid OME-Zarr well metadata.
+        """
+        return _from_zarr_v3(group, cls, WellAttrs)