From 311e78caea9d9350d2f75dfccce13228740a4192 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sun, 5 Jan 2025 18:00:44 +0000
Subject: [PATCH] Add a base for group objects

---
 src/ome_zarr_models/__init__.py        | 37 ++++++++++++++++++++++++++
 src/ome_zarr_models/base.py            | 18 +++++++++++++
 src/ome_zarr_models/v04/base.py        | 16 +++++++++++
 src/ome_zarr_models/v04/hcs.py         |  3 ++-
 src/ome_zarr_models/v04/image.py       |  3 ++-
 src/ome_zarr_models/v04/image_label.py |  3 ++-
 src/ome_zarr_models/v04/labels.py      |  3 ++-
 src/ome_zarr_models/v04/well.py        |  3 ++-
 8 files changed, 81 insertions(+), 5 deletions(-)
 create mode 100644 src/ome_zarr_models/v04/base.py

diff --git a/src/ome_zarr_models/__init__.py b/src/ome_zarr_models/__init__.py
index 7db6724..8ef07af 100644
--- a/src/ome_zarr_models/__init__.py
+++ b/src/ome_zarr_models/__init__.py
@@ -1,6 +1,43 @@
 from importlib.metadata import PackageNotFoundError, version
 
+import zarr
+
+import ome_zarr_models.v04.hcs
+import ome_zarr_models.v04.image
+import ome_zarr_models.v04.image_label
+import ome_zarr_models.v04.labels
+import ome_zarr_models.v04.well
+from ome_zarr_models.base import BaseGroup
+from ome_zarr_models.v04.base import BaseGroupv04
+
 try:
     __version__ = version("ome_zarr_models")
 except PackageNotFoundError:  # pragma: no cover
     __version__ = "uninstalled"
+
+
+_V04_groups: list[type[BaseGroupv04]] = [
+    ome_zarr_models.v04.hcs.HCS,
+    ome_zarr_models.v04.image_label.ImageLabel,
+    ome_zarr_models.v04.image.Image,
+    ome_zarr_models.v04.labels.Labels,
+    ome_zarr_models.v04.well.Well,
+]
+
+
+def load_ome_zarr_group(group: zarr.Group) -> BaseGroup:
+    """
+    Create an ome-zarr-models object from an existing OME-Zarr group.
+
+    This function will 'guess' which type of OME-Zarr data exists by
+    trying to validate each group metadata definition against your data.
+    If validation is successful, that data class is returned without
+    trying any more.
+    """
+    for group_cls in _V04_groups:
+        try:
+            return group_cls.from_zarr(group)
+        except Exception:
+            continue
+
+    raise RuntimeError(f"Could not find any matches for group {group}")
diff --git a/src/ome_zarr_models/base.py b/src/ome_zarr_models/base.py
index c7559be..398ca7f 100644
--- a/src/ome_zarr_models/base.py
+++ b/src/ome_zarr_models/base.py
@@ -1,4 +1,8 @@
+from abc import ABC, abstractmethod
+from typing import Literal
+
 from pydantic import BaseModel, ConfigDict
+from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 
 class BaseAttrs(BaseModel):
@@ -17,3 +21,17 @@ class BaseAttrs(BaseModel):
         populate_by_name=True,
         frozen=True,
     )
+
+
+class BaseGroup(GroupSpec[BaseAttrs, ArraySpec | GroupSpec], ABC):
+    """
+    Base class for all OME-Zarr groups.
+    """
+
+    @property
+    @abstractmethod
+    def ome_zarr_version(self) -> Literal["0.4"]:
+        """
+        Version of the OME-Zarr specification that this group corresponds to.
+        """
+        ...
diff --git a/src/ome_zarr_models/v04/base.py b/src/ome_zarr_models/v04/base.py
new file mode 100644
index 0000000..9567013
--- /dev/null
+++ b/src/ome_zarr_models/v04/base.py
@@ -0,0 +1,16 @@
+from typing import Literal
+
+from ome_zarr_models.base import BaseGroup
+
+
+class BaseGroupv04(BaseGroup):
+    """
+    Base class for all v0.4 OME-Zarr groups.
+    """
+
+    @property
+    def ome_zarr_version(self) -> Literal["0.4"]:
+        """
+        OME-Zarr version.
+        """
+        return "0.4"
diff --git a/src/ome_zarr_models/v04/hcs.py b/src/ome_zarr_models/v04/hcs.py
index dcff3bb..3270abe 100644
--- a/src/ome_zarr_models/v04/hcs.py
+++ b/src/ome_zarr_models/v04/hcs.py
@@ -5,6 +5,7 @@
 from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 from ome_zarr_models.base import BaseAttrs
+from ome_zarr_models.v04.base import BaseGroupv04
 from ome_zarr_models.v04.plate import Plate
 from ome_zarr_models.v04.well import Well
 
@@ -19,7 +20,7 @@ class HCSAttrs(BaseAttrs):
     plate: Plate
 
 
-class HCS(GroupSpec[HCSAttrs, ArraySpec | GroupSpec]):  # type: ignore[misc]
+class HCS(GroupSpec[HCSAttrs, ArraySpec | GroupSpec], BaseGroupv04):  # type: ignore[misc]
     """
     An OME-zarr high-content screening (HCS) dataset representing a single plate.
     """
diff --git a/src/ome_zarr_models/v04/image.py b/src/ome_zarr_models/v04/image.py
index c1971d3..8ab9ae1 100644
--- a/src/ome_zarr_models/v04/image.py
+++ b/src/ome_zarr_models/v04/image.py
@@ -8,6 +8,7 @@
 
 from ome_zarr_models._utils import get_store_path
 from ome_zarr_models.base import BaseAttrs
+from ome_zarr_models.v04.base import BaseGroupv04
 from ome_zarr_models.v04.labels import Labels
 from ome_zarr_models.v04.multiscales import Multiscale
 from ome_zarr_models.v04.omero import Omero
@@ -74,7 +75,7 @@ class ImageAttrs(BaseAttrs):
     omero: Omero | None = None
 
 
-class Image(GroupSpec[ImageAttrs, ArraySpec | GroupSpec]):  # type: ignore[misc]
+class Image(GroupSpec[ImageAttrs, ArraySpec | GroupSpec], BaseGroupv04):  # type: ignore[misc]
     """
     An OME-zarr multiscale dataset.
     """
diff --git a/src/ome_zarr_models/v04/image_label.py b/src/ome_zarr_models/v04/image_label.py
index 1f6e4a1..419a06a 100644
--- a/src/ome_zarr_models/v04/image_label.py
+++ b/src/ome_zarr_models/v04/image_label.py
@@ -11,6 +11,7 @@
 from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 from ome_zarr_models.base import BaseAttrs
+from ome_zarr_models.v04.base import BaseGroupv04
 from ome_zarr_models.v04.image import Image, _check_arrays_compatible
 from ome_zarr_models.v04.image_label_types import (
     Label,
@@ -36,7 +37,7 @@ class ImageLabelAttrs(BaseAttrs):
     multiscales: list[Multiscale]
 
 
-class ImageLabel(GroupSpec[ImageLabelAttrs, ArraySpec | GroupSpec]):  # type: ignore[misc]
+class ImageLabel(GroupSpec[ImageLabelAttrs, ArraySpec | GroupSpec], BaseGroupv04):  # type: ignore[misc]
     """
     An image label dataset.
     """
diff --git a/src/ome_zarr_models/v04/labels.py b/src/ome_zarr_models/v04/labels.py
index cbbf37e..3e00862 100644
--- a/src/ome_zarr_models/v04/labels.py
+++ b/src/ome_zarr_models/v04/labels.py
@@ -4,6 +4,7 @@
 from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 from ome_zarr_models.base import BaseAttrs
+from ome_zarr_models.v04.base import BaseGroupv04
 
 __all__ = ["Labels", "LabelsAttrs"]
 
@@ -18,7 +19,7 @@ class LabelsAttrs(BaseAttrs):
     )
 
 
-class Labels(GroupSpec[LabelsAttrs, ArraySpec | GroupSpec]):  # type: ignore[misc]
+class Labels(GroupSpec[LabelsAttrs, ArraySpec | GroupSpec], BaseGroupv04):  # type: ignore[misc]
     """
     An OME-zarr labels dataset.
     """
diff --git a/src/ome_zarr_models/v04/well.py b/src/ome_zarr_models/v04/well.py
index a703106..82678fa 100644
--- a/src/ome_zarr_models/v04/well.py
+++ b/src/ome_zarr_models/v04/well.py
@@ -7,6 +7,7 @@
 from pydantic_zarr.v2 import ArraySpec, GroupSpec
 
 from ome_zarr_models.base import BaseAttrs
+from ome_zarr_models.v04.base import BaseGroupv04
 from ome_zarr_models.v04.image import Image
 from ome_zarr_models.v04.well_types import WellMeta
 
@@ -22,7 +23,7 @@ class WellAttrs(BaseAttrs):
     well: WellMeta
 
 
-class Well(GroupSpec[WellAttrs, ArraySpec | GroupSpec]):  # type: ignore[misc]
+class Well(GroupSpec[WellAttrs, ArraySpec | GroupSpec], BaseGroupv04):  # type: ignore[misc]
     def get_image(self, i: int) -> Image:
         """
         Get a single image from this well.