refactor(ngff): extract get_ome_attrs utility for version-independent OME metadata access

Replace inline v0.4/v0.5 metadata branching with get_ome_attrs() from core.compat. Used in NGFFNode.maybe_wrapped_ome_attrs, Plate.first_position, and tests.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>

Author: srivarra

src/iohub/core/__init__.py

@@ -1,7 +1,7 @@
 """Core zarr implementation abstraction for iohub."""
 
 from iohub.core.arrays import NGFFArray
-from iohub.core.compat import ngff_version_for_format, zarr_format_for_version
+from iohub.core.compat import get_ome_attrs, ngff_version_for_format, zarr_format_for_version
 from iohub.core.config import (
     CompressorConfig,
     ImplementationConfig,

src/iohub/core/compat.py

@@ -16,6 +16,16 @@ def zarr_format_for_version(version: NGFFVersion) -> ZarrFormat:
         raise ValueError(f"Unknown NGFF version: {version!r}. Supported: {list(NGFF_TO_ZARR_FORMAT)}") from err
 
 
+def get_ome_attrs(attrs) -> dict:
+    """Extract OME metadata dict from zarr attrs, regardless of NGFF version.
+
+    v0.5 stores wrap metadata under an ``"ome"`` key; v0.4 stores
+    place it flat in ``.zattrs``.  This function returns the OME
+    metadata dict in both cases.
+    """
+    return attrs.get("ome") or dict(attrs)
+
+
 def ngff_version_for_format(zarr_format: ZarrFormat) -> NGFFVersion:
     """Map zarr format integer to NGFF version string."""
     try:

src/iohub/ngff/nodes.py

@@ -26,6 +26,7 @@
 from pydantic import ValidationError
 
 from iohub.core import ArraySpec, NGFFArray, get_implementation
+from iohub.core.compat import get_ome_attrs
 from iohub.core.config import ImplementationConfig
 from iohub.core.errors import StoreOpenError
 from iohub.core.protocol import ZarrImplementation
@@ -184,7 +185,7 @@ def zattrs(self):
     @property
     def maybe_wrapped_ome_attrs(self):
         """Container of OME metadata attributes."""
-        return self.zattrs.get("ome") or self.zattrs
+        return get_ome_attrs(self.zattrs)
 
     @property
     def version(self) -> Literal["0.4", "0.5"]:
@@ -2486,7 +2487,7 @@ def _first_pos(self):
         try:
             well_path = self.metadata.wells[0].path
             well_grp = self.zgroup[well_path]
-            attrs = well_grp.attrs.get("ome") or dict(well_grp.attrs)
+            attrs = get_ome_attrs(well_grp.attrs)
             pos_name = attrs["well"]["images"][0]["path"]
             return Position(
                 group=well_grp[pos_name],

tests/ngff/test_ngff.py

@@ -25,6 +25,7 @@
 if TYPE_CHECKING:
     from _typeshed import StrPath
 
+from iohub.core.compat import get_ome_attrs
 from iohub.core.utils import pad_shape
 from iohub.ngff.models import TO_DICT_SETTINGS
 from iohub.ngff.nodes import (
@@ -825,10 +826,7 @@ def test_set_transform_fov(ch_shape_dtype, arr_name, version):
             assert dataset.metadata.multiscales[0].coordinate_transformations == transform
         # read data with plain zarr
         group = zarr.open(store_path)
-        if version == "0.4":
-            maybe_ome = group.attrs
-        elif version == "0.5":
-            maybe_ome = group.attrs["ome"]
+        maybe_ome = get_ome_attrs(group.attrs)
         assert maybe_ome["multiscales"][0]["coordinateTransformations"] == [
             translate.model_dump(**TO_DICT_SETTINGS) for translate in transform
         ]
@@ -1157,7 +1155,7 @@ def test_create_well(row_names: list[str], col_names: list[str]):
         for row_name in row_names:
             for col_name in col_names:
                 dataset.create_well(row_name, col_name)
-        plate_meta = dataset.zattrs.get("ome", dataset.zattrs)["plate"]
+        plate_meta = get_ome_attrs(dataset.zattrs)["plate"]
         assert [c["name"] for c in plate_meta["columns"]] == col_names
         assert [r["name"] for r in plate_meta["rows"]] == row_names
 
@@ -1206,7 +1204,7 @@ def test_create_position(row, col, pos, version):
             version=version,
         )
         _ = dataset.create_position(row_name=row, col_name=col, pos_name=pos)
-        ome = dict(dataset.zgroup.attrs) if version == "0.4" else dataset.zgroup.attrs["ome"]
+        ome = get_ome_attrs(dataset.zgroup.attrs)
         assert [c["name"] for c in ome["plate"]["columns"]] == [col]
         assert [r["name"] for r in ome["plate"]["rows"]] == [row]
         assert (store_path / row / col / pos).is_dir()
@@ -1257,7 +1255,7 @@ def test_create_positions(tmp_path, version):
 
     # Collect positions and compare those
 
-    get_metadata = lambda x: dict(x.zgroup.attrs["ome"])
+    get_metadata = lambda x: get_ome_attrs(x.zgroup.attrs)
 
     single_plate_metadata = get_metadata(single)
     batched_plate_metadata = get_metadata(batched)
@@ -1316,7 +1314,7 @@ def test_create_positions_with_tuple_variations(tmp_path, version):
     batched.create_positions(positions)
 
     # Verify metadata matches
-    get_metadata = lambda x: dict(x.zgroup.attrs["ome"])
+    get_metadata = lambda x: get_ome_attrs(x.zgroup.attrs)
 
     single_meta = get_metadata(single)
     batched_meta = get_metadata(batched)