ensure we can override and add test

Author: d33bs

src/ome_arrow/core.py

@@ -59,6 +59,7 @@ def __init__(
         tcz: Tuple[int, int, int] = (0, 0, 0),
         column_name: str = "ome_arrow",
         row_index: int = 0,
+        image_type: str | None = None,
     ) -> None:
         """
         Construct an OMEArrow from:
@@ -71,6 +72,7 @@ def __init__(
             with from_numpy defaults)
         - a dict already matching the OME-Arrow schema
         - a pa.StructScalar already typed to OME_ARROW_STRUCT
+        - optionally override/set image_type metadata on ingest
         """
 
         # set the tcz for viewing
@@ -83,6 +85,7 @@ def __init__(
                 default_dim_for_unspecified="C",
                 map_series_to="T",
                 clamp_to_uint16=True,
+                image_type=image_type,
             )
 
         # --- 2) String path/URL: OME-Zarr / OME-Parquet / OME-TIFF ---------------
@@ -98,6 +101,14 @@ def __init__(
                 or (path.exists() and path.is_dir() and path.suffix.lower() == ".zarr")
             ):
                 self.data = from_ome_zarr(s)
+                if image_type is not None:
+                    self.data = pa.scalar(
+                        {
+                            **self.data.as_py(),
+                            "image_type": str(image_type),
+                        },
+                        type=OME_ARROW_STRUCT,
+                    )
 
             # OME-Parquet
             elif s.lower().endswith((".parquet", ".pq")) or path.suffix.lower() in {
@@ -107,18 +118,42 @@ def __init__(
                 self.data = from_ome_parquet(
                     s, column_name=column_name, row_index=row_index
                 )
+                if image_type is not None:
+                    self.data = pa.scalar(
+                        {
+                            **self.data.as_py(),
+                            "image_type": str(image_type),
+                        },
+                        type=OME_ARROW_STRUCT,
+                    )
 
             # Vortex
             elif s.lower().endswith(".vortex") or path.suffix.lower() == ".vortex":
                 self.data = from_ome_vortex(
                     s, column_name=column_name, row_index=row_index
                 )
+                if image_type is not None:
+                    self.data = pa.scalar(
+                        {
+                            **self.data.as_py(),
+                            "image_type": str(image_type),
+                        },
+                        type=OME_ARROW_STRUCT,
+                    )
 
             # TIFF
             elif path.suffix.lower() in {".tif", ".tiff"} or s.lower().endswith(
                 (".tif", ".tiff")
             ):
                 self.data = from_tiff(s)
+                if image_type is not None:
+                    self.data = pa.scalar(
+                        {
+                            **self.data.as_py(),
+                            "image_type": str(image_type),
+                        },
+                        type=OME_ARROW_STRUCT,
+                    )
 
             elif path.exists() and path.is_dir():
                 raise ValueError(
@@ -140,15 +175,33 @@ def __init__(
             # Uses from_numpy defaults: dim_order="TCZYX", clamp_to_uint16=True, etc.
             # If the array is YX/ZYX/CYX/etc.,
             # from_numpy will expand/reorder accordingly.
-            self.data = from_numpy(data)
+            self.data = from_numpy(data, image_type=image_type)
 
         # --- 4) Already-typed Arrow scalar ---------------------------------------
         elif isinstance(data, pa.StructScalar):
-            self.data = data
+            if image_type is None:
+                self.data = data
+            else:
+                self.data = pa.scalar(
+                    {
+                        **data.as_py(),
+                        "image_type": str(image_type),
+                    },
+                    type=OME_ARROW_STRUCT,
+                )
 
         # --- 5) Plain dict matching the schema -----------------------------------
         elif isinstance(data, dict):
-            self.data = pa.scalar(data, type=OME_ARROW_STRUCT)
+            if image_type is None:
+                self.data = pa.scalar(data, type=OME_ARROW_STRUCT)
+            else:
+                self.data = pa.scalar(
+                    {
+                        **data,
+                        "image_type": str(image_type),
+                    },
+                    type=OME_ARROW_STRUCT,
+                )
 
         # --- otherwise ------------------------------------------------------------
         else:

tests/test_core.py

@@ -369,6 +369,18 @@ def test_vortex_roundtrip(tmp_path: pathlib.Path) -> None:
     assert reloaded.info() == oa.info()
 
 
+def test_parquet_roundtrip_preserves_image_type(tmp_path: pathlib.Path) -> None:
+    """Ensure image_type round-trips through OME-Parquet."""
+    arr = np.arange(16, dtype=np.uint16).reshape(1, 1, 1, 4, 4)
+    oa = OMEArrow(arr, image_type="label")
+    out = tmp_path / "example.ome.parquet"
+
+    oa.export(how="omeparquet", out=str(out))
+    reloaded = OMEArrow(str(out))
+
+    assert reloaded.data.as_py()["image_type"] == "label"
+
+
 def test_vortex_custom_column_name(tmp_path: pathlib.Path) -> None:
     """Ensure custom Vortex column names are preserved on round-trip."""
     pytest.importorskip(