Add interface for getting wells in a HCS

src/ome_zarr_models/v04/hcs.py

@@ -1,7 +1,10 @@
+from collections.abc import Generator
+
 from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 from ome_zarr_models.base import Base
 from ome_zarr_models.v04.plate import Plate
+from ome_zarr_models.v04.well import WellGroup
 
 __all__ = ["HCSAttrs"]
 
@@ -18,3 +21,28 @@ class HCS(GroupSpec[HCSAttrs, ArraySpec | GroupSpec]):
     """
     An OME-zarr high-content screening (HCS) dataset representing a single plate.
     """
+
+    def get_well_group(self, i: int) -> WellGroup:
+        """
+        Get a single well group.
+        """
+        well = self.attributes.plate.wells[i]
+        well_path = well.path
+        well_path_parts = well_path.split("/")
+        group = self
+        for part in well_path_parts:
+            group = group.members[part]
+
+        return WellGroup(attributes=group.attributes, members=group.members)
+
+    @property
+    def n_wells(self) -> int:
+        """
+        Number of wells.
+        """
+        return len(self.attributes.plate.wells)
+
+    @property
+    def well_groups(self) -> Generator[WellGroup, None, None]:
+        for i in range(self.n_wells):
+            yield self.get_well_group(i)

src/ome_zarr_models/v04/well.py

@@ -6,11 +6,13 @@
 from typing import Annotated, Literal
 
 from pydantic import AfterValidator, Field
+from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 from ome_zarr_models._utils import _AlphaNumericConstraint, _unique_items_validator
 from ome_zarr_models.base import Base
 
-__all__ = ["Well", "WellImage"]
+# WellGroup is defined one level higher
+__all__ = ["Well", "WellAttrs", "WellImage"]
 
 
 class WellImage(Base):
@@ -58,3 +60,15 @@ def get_acquisition_paths(self) -> dict[int, list[str]]:
                 )
             acquisition_dict[image.acquisition].append(image.path)
         return dict(acquisition_dict)
+
+
+class WellAttrs(Base):
+    """
+    Attributes for a well group.
+    """
+
+    well: Well
+
+
+class WellGroup(GroupSpec[WellAttrs, ArraySpec | GroupSpec]):
+    pass

tests/v04/test_hcs.py

@@ -5,6 +5,7 @@
 from ome_zarr_models.v04 import HCS
 from ome_zarr_models.v04.hcs import HCSAttrs
 from ome_zarr_models.v04.plate import Acquisition, Column, Plate, Row, WellInPlate
+from ome_zarr_models.v04.well import Well, WellImage
 
 
 def test_example_hcs() -> None:
@@ -30,3 +31,9 @@ def test_example_hcs() -> None:
             wells=[WellInPlate(path="B/03", rowIndex=0, columnIndex=0)],
         )
     )
+
+    well_groups = list(hcs.well_groups)
+    assert len(well_groups) == 1
+    assert well_groups[0].attributes.well == Well(
+        images=[WellImage(path="0", acquisition=None)], version="0.4"
+    )