Clean up types/inheritance for v05 (#138)

ome-zarr-models · Jan 30, 2025 · 9c610fe · 9c610fe
1 parent 681c241
commit 9c610fe
Show file tree

Hide file tree

Showing 11 changed files with 19 additions and 43 deletions.
diff --git a/src/ome_zarr_models/_v05/base.py b/src/ome_zarr_models/_v05/base.py
@@ -1,20 +1,20 @@
-from typing import Literal
-
-from pydantic_zarr.v2 import ArraySpec, GroupSpec
+from typing import Generic, Literal, TypeVar
 
 from ome_zarr_models.base import BaseAttrs, BaseGroup
 
+T = TypeVar("T", bound=BaseAttrs)
+
 
-class BaseOMEAttrs(BaseAttrs):
+class BaseOMEAttrs(BaseAttrs, Generic[T]):
     """
     Base class for all OME attributes.
     """
 
     version: Literal["0.5"] = "0.5"
-    ome: BaseAttrs
+    ome: T
 
 
-class BaseGroupv05(GroupSpec[BaseOMEAttrs, ArraySpec | GroupSpec], BaseGroup):  # type: ignore[misc]
+class BaseGroupv05(BaseGroup):
     """
     Base class for all v0.5 OME-Zarr groups.
     """
@@ -25,10 +25,3 @@ def ome_zarr_version(self) -> Literal["0.5"]:
         OME-Zarr version.
         """
         return "0.5"
-
-    @property
-    def ome_attributes(self) -> BaseAttrs:
-        """
-        OME-Zarr attributes.
-        """
-        return self.attributes.ome  # type: ignore[no-any-return]
diff --git a/src/ome_zarr_models/_v05/hcs.py b/src/ome_zarr_models/_v05/hcs.py
@@ -15,11 +15,7 @@ class HCSAttrs(BaseAttrs):
     plate: Plate
 
 
-class OMEHCSAttrs(BaseOMEAttrs):
-    ome: HCSAttrs
-
-
-class HCS(GroupSpec[OMEHCSAttrs, ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
+class HCS(GroupSpec[BaseOMEAttrs[HCSAttrs], ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
     """
     An OME-Zarr high content screening (HCS) dataset.
     """
diff --git a/src/ome_zarr_models/_v05/image.py b/src/ome_zarr_models/_v05/image.py
@@ -6,11 +6,7 @@
 __all__ = ["Image", "ImageAttrs"]
 
 
-class OMEImageAttrs(BaseOMEAttrs):
-    ome: ImageAttrs
-
-
-class Image(GroupSpec[OMEImageAttrs, ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
+class Image(GroupSpec[BaseOMEAttrs[ImageAttrs], ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
     """
     An OME-Zarr image dataset.
     """
diff --git a/src/ome_zarr_models/_v05/image_label.py b/src/ome_zarr_models/_v05/image_label.py
@@ -18,11 +18,10 @@ class ImageLabelAttrs(BaseAttrs):
     multiscales: list[Multiscale]
 
 
-class OMEImageLabelAttrs(BaseOMEAttrs):
-    ome: ImageLabelAttrs
-
-
-class ImageLabel(GroupSpec[OMEImageLabelAttrs, ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
+class ImageLabel(
+    GroupSpec[BaseOMEAttrs[ImageLabelAttrs], ArraySpec | GroupSpec],  # type: ignore[misc]
+    BaseGroupv05,
+):
     """
     An OME-Zarr image label dataset.
     """
diff --git a/src/ome_zarr_models/_v05/labels.py b/src/ome_zarr_models/_v05/labels.py
@@ -6,11 +6,7 @@
 __all__ = ["Labels", "LabelsAttrs"]
 
 
-class OMELabelsAttrs(BaseOMEAttrs):
-    ome: LabelsAttrs
-
-
-class Labels(GroupSpec[OMELabelsAttrs, ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
+class Labels(GroupSpec[BaseOMEAttrs[LabelsAttrs], ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
     """
     An OME-Zarr labels dataset.
     """
diff --git a/src/ome_zarr_models/_v05/well.py b/src/ome_zarr_models/_v05/well.py
@@ -6,11 +6,7 @@
 __all__ = ["Well", "WellAttrs"]
 
 
-class OMEWellAttrs(BaseOMEAttrs):
-    ome: WellAttrs
-
-
-class Well(GroupSpec[OMEWellAttrs, ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
+class Well(GroupSpec[BaseOMEAttrs[WellAttrs], ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
     """
     An OME-Zarr well dataset.
     """
diff --git a/tests/v05/test_hcs.py b/tests/v05/test_hcs.py
@@ -6,7 +6,7 @@
 def test_hcs() -> None:
     zarr_group = json_to_zarr_group(json_fname="hcs_example.json")
     ome_group = HCS.from_zarr(zarr_group)
-    assert ome_group.ome_attributes == HCSAttrs(
+    assert ome_group.attributes.ome == HCSAttrs(
         plate=Plate(
             acquisitions=[
                 Acquisition(

diff --git a/tests/v05/test_image.py b/tests/v05/test_image.py
@@ -8,7 +8,7 @@
 def test_image() -> None:
     zarr_group = json_to_zarr_group(json_fname="image_example.json")
     ome_group = Image.from_zarr(zarr_group)
-    assert ome_group.ome_attributes == ImageAttrs(
+    assert ome_group.attributes.ome == ImageAttrs(
         multiscales=[
             Multiscale(
                 axes=[

diff --git a/tests/v05/test_image_label.py b/tests/v05/test_image_label.py
@@ -9,7 +9,7 @@
 def test_image_label() -> None:
     zarr_group = json_to_zarr_group(json_fname="image_label_example.json")
     ome_group = ImageLabel.from_zarr(zarr_group)
-    assert ome_group.ome_attributes == ImageLabelAttrs(
+    assert ome_group.attributes.ome == ImageLabelAttrs(
         image_label=Label(
             colors=(
                 Color(label_value=0, rgba=(0, 0, 128, 128)),

diff --git a/tests/v05/test_labels.py b/tests/v05/test_labels.py
@@ -5,6 +5,6 @@
 def test_labels() -> None:
     zarr_group = json_to_zarr_group(json_fname="labels_example.json")
     ome_group = Labels.from_zarr(zarr_group)
-    assert ome_group.ome_attributes == LabelsAttrs(
+    assert ome_group.attributes.ome == LabelsAttrs(
         labels=["cell_space_segmentation"], version="0.5"
     )
diff --git a/tests/v05/test_well.py b/tests/v05/test_well.py
@@ -6,7 +6,7 @@
 def test_well() -> None:
     zarr_group = json_to_zarr_group(json_fname="well_example.json")
     ome_group = Well.from_zarr(zarr_group)
-    assert ome_group.ome_attributes == WellAttrs(
+    assert ome_group.attributes.ome == WellAttrs(
         well=WellMeta(
             images=[
                 WellImage(path="0", acquisition=1),