Handle unparsed streaming artifact metadata

Author: trissim

src/polystore/__init__.py

@@ -26,10 +26,10 @@
     get_backend,
 )
 from .constants import Backend, MemoryType, TransportMode
-from .disk import DiskStorageBackend
+from .disk import DiskBackend, DiskStorageBackend
 from .filemanager import FileManager
 from .formats import FileFormat, DEFAULT_IMAGE_EXTENSIONS
-from .memory import MemoryStorageBackend
+from .memory import MemoryBackend, MemoryStorageBackend
 from .metadata_writer import (
     AtomicMetadataWriter,
     MetadataWriteError,
@@ -76,7 +76,9 @@
     "register_cleanup_callback",
     "STORAGE_BACKENDS",
     "DiskStorageBackend",
+    "DiskBackend",
     "MemoryStorageBackend",
+    "MemoryBackend",
     "FileManager",
     "file_lock",
     "atomic_write_json",

src/polystore/disk.py

@@ -834,3 +834,6 @@ def _save_rois(self, rois: List, output_path: Path, images_dir: str = None, **kw
 
         logger.info(f"Saved {roi_count} ROIs to .roi.zip archive: {output_path}")
         return str(output_path)
+
+
+DiskBackend = DiskStorageBackend

src/polystore/memory.py

@@ -657,3 +657,6 @@ def __init__(self, target: str):
 
     def __repr__(self):
         return f"<MemorySymlink → {self.target}>"
+
+
+MemoryBackend = MemoryStorageBackend

src/polystore/streaming/_streaming_backend.py

@@ -9,8 +9,9 @@
 import os
 import time
 import uuid
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, List, Set, Union
+from typing import Any, Callable, List, Mapping, Set, Union
 import numpy as np
 
 from ..base import DataSink
@@ -24,6 +25,27 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class StreamingComponentMetadata:
+    """Message metadata for one streamed item."""
+
+    parsed_filename_metadata: Mapping[str, Any] | None
+    source: str
+
+    def to_payload(self) -> dict[str, Any]:
+        if self.parsed_filename_metadata is None:
+            metadata: dict[str, Any] = {}
+        elif isinstance(self.parsed_filename_metadata, Mapping):
+            metadata = dict(self.parsed_filename_metadata)
+        else:
+            raise TypeError(
+                "Streaming filename parser must return a mapping or None, "
+                f"got {type(self.parsed_filename_metadata).__name__}."
+            )
+        metadata["source"] = self.source
+        return metadata
+
+
 class StreamingBackend(DataSink):
     """
     Abstract base class for ZeroMQ-based streaming backends.
@@ -165,12 +187,10 @@ def _parse_component_metadata(self, file_path: Union[str, Path], microscope_hand
             Component metadata dict with source added
         """
         filename = os.path.basename(str(file_path))
-        component_metadata = microscope_handler.parser.parse_filename(filename)
-
-        # Add pre-built source value directly
-        component_metadata['source'] = source
-
-        return component_metadata
+        return StreamingComponentMetadata(
+            microscope_handler.parser.parse_filename(filename),
+            source,
+        ).to_payload()
 
     def _detect_data_type(self, data: Any):
         """

tests/test_streaming_metadata.py

@@ -0,0 +1,78 @@
+from types import SimpleNamespace
+
+import pytest
+
+from polystore.streaming._streaming_backend import StreamingBackend
+
+
+class MetadataProbeStreamingBackend(StreamingBackend):
+    VIEWER_TYPE = "probe"
+    SHM_PREFIX = "probe_"
+
+    def save_batch(self, data_list, file_paths, **kwargs):
+        raise NotImplementedError
+
+
+def test_streaming_component_metadata_accepts_unparsed_artifact_filename() -> None:
+    backend = MetadataProbeStreamingBackend()
+    microscope_handler = SimpleNamespace(
+        parser=SimpleNamespace(parse_filename=lambda _filename: None)
+    )
+
+    metadata = backend._parse_component_metadata(
+        "A01_s001_w1_z001_t001_Nuclei_step3_rois.roi.zip",
+        microscope_handler,
+        source="IdentifyPrimaryObjects",
+    )
+
+    assert metadata == {"source": "IdentifyPrimaryObjects"}
+
+
+def test_streaming_batch_items_accept_unparsed_artifact_filename() -> None:
+    backend = MetadataProbeStreamingBackend()
+    microscope_handler = SimpleNamespace(
+        parser=SimpleNamespace(parse_filename=lambda _filename: None)
+    )
+
+    batch_images, image_ids = backend._prepare_batch_items(
+        [object()],
+        ["A01_s001_w1_z001_t001_Nuclei_step3_rois.roi.zip"],
+        microscope_handler,
+        "IdentifyPrimaryObjects",
+        lambda _data, _path, _data_type: ({"payload": "ok"}, "image"),
+    )
+
+    assert len(image_ids) == 1
+    assert batch_images[0]["metadata"] == {"source": "IdentifyPrimaryObjects"}
+    assert batch_images[0]["payload"] == "ok"
+
+
+def test_streaming_component_metadata_preserves_parsed_filename_fields() -> None:
+    backend = MetadataProbeStreamingBackend()
+    microscope_handler = SimpleNamespace(
+        parser=SimpleNamespace(
+            parse_filename=lambda _filename: {"well": "A01", "channel": 1}
+        )
+    )
+
+    metadata = backend._parse_component_metadata(
+        "A01_s001_w1_z001_t001.TIF",
+        microscope_handler,
+        source="Crop",
+    )
+
+    assert metadata == {"well": "A01", "channel": 1, "source": "Crop"}
+
+
+def test_streaming_component_metadata_rejects_invalid_parser_result() -> None:
+    backend = MetadataProbeStreamingBackend()
+    microscope_handler = SimpleNamespace(
+        parser=SimpleNamespace(parse_filename=lambda _filename: ["not", "metadata"])
+    )
+
+    with pytest.raises(TypeError, match="mapping or None"):
+        backend._parse_component_metadata(
+            "A01_s001_w1_z001_t001.TIF",
+            microscope_handler,
+            source="Crop",
+        )