From 4959a165e5ae2d6c5ab3149138c132fd534ef46d Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sat, 11 Jan 2025 17:54:41 +0000
Subject: [PATCH] Add labels to v05

---
 src/ome_zarr_models/base.py        |  2 +-
 src/ome_zarr_models/v05/base.py    | 34 ++++++++++++++++++++++++++++++
 src/ome_zarr_models/v05/labels.py  | 16 ++++++++++++++
 tests/v04/__init__.py              |  0
 tests/v05/__init__.py              |  0
 tests/v05/conftest.py              | 21 ++++++++++++++++++
 tests/v05/data/labels_example.json |  6 ++++++
 tests/v05/test_labels.py           | 11 ++++++++++
 8 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 src/ome_zarr_models/v05/base.py
 create mode 100644 src/ome_zarr_models/v05/labels.py
 create mode 100644 tests/v04/__init__.py
 create mode 100644 tests/v05/__init__.py
 create mode 100644 tests/v05/conftest.py
 create mode 100644 tests/v05/data/labels_example.json
 create mode 100644 tests/v05/test_labels.py

diff --git a/src/ome_zarr_models/base.py b/src/ome_zarr_models/base.py
index 1d5aee6..2bdec78 100644
--- a/src/ome_zarr_models/base.py
+++ b/src/ome_zarr_models/base.py
@@ -30,7 +30,7 @@ class BaseGroup(GroupSpec[BaseAttrs, ArraySpec | GroupSpec], ABC):  # type: igno
 
     @property
     @abstractmethod
-    def ome_zarr_version(self) -> Literal["0.4"]:
+    def ome_zarr_version(self) -> Literal["0.4", "0.5"]:
         """
         Version of the OME-Zarr specification that this group corresponds to.
         """
diff --git a/src/ome_zarr_models/v05/base.py b/src/ome_zarr_models/v05/base.py
new file mode 100644
index 0000000..4f3b686
--- /dev/null
+++ b/src/ome_zarr_models/v05/base.py
@@ -0,0 +1,34 @@
+from typing import Literal
+
+from pydantic_zarr.v2 import ArraySpec, GroupSpec
+
+from ome_zarr_models.base import BaseAttrs, BaseGroup
+
+
+class BaseOMEAttrs(BaseAttrs):
+    """
+    Base class for all OME attributes.
+    """
+
+    version: Literal["0.5"] = "0.5"
+    ome: BaseAttrs
+
+
+class BaseGroupv05(GroupSpec[BaseOMEAttrs, ArraySpec | GroupSpec], BaseGroup):  # type: ignore[misc]
+    """
+    Base class for all v0.5 OME-Zarr groups.
+    """
+
+    @property
+    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/labels.py b/src/ome_zarr_models/v05/labels.py
new file mode 100644
index 0000000..23f9fe0
--- /dev/null
+++ b/src/ome_zarr_models/v05/labels.py
@@ -0,0 +1,16 @@
+from pydantic_zarr.v2 import ArraySpec, GroupSpec
+
+from ome_zarr_models.v04.labels import LabelsAttrs
+from ome_zarr_models.v05.base import BaseGroupv05, BaseOMEAttrs
+
+__all__ = ["Labels", "LabelsAttrs"]
+
+
+class OMELabelsAttrs(BaseOMEAttrs):
+    ome: LabelsAttrs
+
+
+class Labels(GroupSpec[OMELabelsAttrs, ArraySpec | GroupSpec], BaseGroupv05):  # type: ignore[misc]
+    """
+    An OME-Zarr labels dataset.
+    """
diff --git a/tests/v04/__init__.py b/tests/v04/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/v05/__init__.py b/tests/v05/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/v05/conftest.py b/tests/v05/conftest.py
new file mode 100644
index 0000000..b41de64
--- /dev/null
+++ b/tests/v05/conftest.py
@@ -0,0 +1,21 @@
+import json
+from pathlib import Path
+from typing import TypeVar
+
+import zarr
+
+from ome_zarr_models.base import BaseAttrs
+
+T = TypeVar("T", bound=BaseAttrs)
+
+
+def json_to_zarr_group(*, json_fname: str) -> zarr.Group:
+    """
+    Create an empty Zarr group, and set attributes from a JSON file.
+    """
+    group = zarr.open_group(store=zarr.MemoryStore())
+    with open(Path(__file__).parent / "data" / json_fname) as f:
+        attrs = json.load(f)
+
+    group.attrs.put(attrs)
+    return group
diff --git a/tests/v05/data/labels_example.json b/tests/v05/data/labels_example.json
new file mode 100644
index 0000000..1506d61
--- /dev/null
+++ b/tests/v05/data/labels_example.json
@@ -0,0 +1,6 @@
+{
+  "ome": {
+    "version": "0.5",
+    "labels": ["cell_space_segmentation"]
+  }
+}
diff --git a/tests/v05/test_labels.py b/tests/v05/test_labels.py
new file mode 100644
index 0000000..a6d9f57
--- /dev/null
+++ b/tests/v05/test_labels.py
@@ -0,0 +1,11 @@
+from ome_zarr_models.v04.labels import LabelsAttrs
+from ome_zarr_models.v05.labels import Labels, OMELabelsAttrs
+from tests.v05.conftest import json_to_zarr_group
+
+
+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(
+        labels=["cell_space_segmentation"], version="0.5"
+    )