pymmcore-plus 0.9.5__tar.gz → 0.10.1__tar.gz

{pymmcore_plus-0.9.5 → pymmcore_plus-0.10.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pymmcore-plus
-Version: 0.9.5
+Version: 0.10.1
 Summary: pymmcore superset providing improved APIs, event handling, and a pure python acquisition engine
 Project-URL: Source, https://github.com/pymmcore-plus/pymmcore-plus
 Project-URL: Tracker, https://github.com/pymmcore-plus/pymmcore-plus/issues
@@ -30,6 +30,7 @@ Requires-Dist: platformdirs>=3.0.0
 Requires-Dist: psygnal>=0.7
 Requires-Dist: pymmcore>=10.7.0.71.0
 Requires-Dist: rich>=10.2.0
+Requires-Dist: tensorstore
 Requires-Dist: typer>=0.4.2
 Requires-Dist: typing-extensions
 Requires-Dist: useq-schema>=0.4.7
@@ -43,6 +44,7 @@ Requires-Dist: mypy; extra == 'dev'
 Requires-Dist: pdbpp; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: ruff; extra == 'dev'
+Requires-Dist: tensorstore-stubs; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: mkdocs-material; extra == 'docs'
 Requires-Dist: mkdocs>=1.4; extra == 'docs'
@@ -70,7 +72,7 @@ Description-Content-Type: text/markdown
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/pymmcore-plus)](https://pypi.org/project/pymmcore-plus)
 [![PyPI](https://img.shields.io/pypi/v/pymmcore-plus.svg?color=green)](https://pypi.org/project/pymmcore-plus)
 [![Conda](https://img.shields.io/conda/vn/conda-forge/pymmcore-plus)](https://anaconda.org/conda-forge/pymmcore-plus)
-[![CI](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/test_and_deploy.yml/badge.svg)](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/test_and_deploy.yml)
+[![CI](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/ci.yml/badge.svg)](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/ci.yml)
 [![docs](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/docs.yml/badge.svg)](https://pymmcore-plus.github.io/pymmcore-plus/)
 [![codecov](https://codecov.io/gh/pymmcore-plus/pymmcore-plus/branch/main/graph/badge.svg)](https://codecov.io/gh/pymmcore-plus/pymmcore-plus)
 [![Benchmarks](https://img.shields.io/endpoint?url=https://codspeed.io/badge.json)](https://codspeed.io/pymmcore-plus/pymmcore-plus)

{pymmcore_plus-0.9.5 → pymmcore_plus-0.10.1}/README.md

@@ -4,7 +4,7 @@
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/pymmcore-plus)](https://pypi.org/project/pymmcore-plus)
 [![PyPI](https://img.shields.io/pypi/v/pymmcore-plus.svg?color=green)](https://pypi.org/project/pymmcore-plus)
 [![Conda](https://img.shields.io/conda/vn/conda-forge/pymmcore-plus)](https://anaconda.org/conda-forge/pymmcore-plus)
-[![CI](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/test_and_deploy.yml/badge.svg)](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/test_and_deploy.yml)
+[![CI](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/ci.yml/badge.svg)](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/ci.yml)
 [![docs](https://github.com/pymmcore-plus/pymmcore-plus/actions/workflows/docs.yml/badge.svg)](https://pymmcore-plus.github.io/pymmcore-plus/)
 [![codecov](https://codecov.io/gh/pymmcore-plus/pymmcore-plus/branch/main/graph/badge.svg)](https://codecov.io/gh/pymmcore-plus/pymmcore-plus)
 [![Benchmarks](https://img.shields.io/endpoint?url=https://codspeed.io/badge.json)](https://codspeed.io/pymmcore-plus/pymmcore-plus)

{pymmcore_plus-0.9.5 → pymmcore_plus-0.10.1}/pyproject.toml

@@ -42,6 +42,7 @@ dependencies = [
     "typing-extensions",      # not actually required at runtime
     "useq-schema >=0.4.7",
     "wrapt >=1.14",
+    "tensorstore",
     # cli requirements included by default for now
     "typer >=0.4.2",
     "rich >=10.2.0",
@@ -64,7 +65,7 @@ test = [
     "zarr >=2.2",
     "xarray",
 ]
-dev = ["ipython", "mypy", "pdbpp", "pre-commit", "ruff"]
+dev = ["ipython", "mypy", "pdbpp", "pre-commit", "ruff", "tensorstore-stubs"]
 docs = [
     "mkdocs >=1.4",
     "mkdocs-material",
@@ -181,3 +182,6 @@ ignore = [
     "tests/**/*",
     "tox.ini",
 ]
+
+[tool.typos.default]
+extend-ignore-identifiers-re = ["(?i)nd2?.*", "(?i)ome"]

{pymmcore_plus-0.9.5 → pymmcore_plus-0.10.1}/src/pymmcore_plus/mda/_engine.py

@@ -176,7 +176,7 @@ class MDAEngine(PMDAEngine):
             # skip if no autofocus device is found
             if not self._mmc.getAutoFocusDevice():
                 logger.warning("No autofocus device found. Cannot execute autofocus.")
-                return ()
+                return
 
             try:
                 # execute hardware autofocus
@@ -191,7 +191,7 @@ class MDAEngine(PMDAEngine):
                 self._z_correction[p_idx] = new_correction + self._z_correction.get(
                     p_idx, 0.0
                 )
-            return ()
+            return
 
         # if the autofocus was engaged at the start of the sequence AND autofocus action
         # did not fail, re-engage it. NOTE: we need to do that AFTER the runner calls
@@ -283,7 +283,7 @@ class MDAEngine(PMDAEngine):
             self._mmc.snapImage()
         except Exception as e:
             logger.warning("Failed to snap image. %s", e)
-            return ()
+            return
         if not event.keep_shutter_open:
             self._mmc.setShutterOpen(False)
         yield ImagePayload(self._mmc.getImage(), event, self.get_frame_metadata())
@@ -326,8 +326,20 @@ class MDAEngine(PMDAEngine):
         """Teardown state of system (hardware, etc.) after `event`."""
         # autoshutter was set at the beginning of the sequence, and this event
         # doesn't want to leave the shutter open.  Re-enable autoshutter.
+        core = self._mmc
         if not event.keep_shutter_open and self._autoshutter_was_set:
-            self._mmc.setAutoShutter(True)
+            core.setAutoShutter(True)
+        # FIXME: this may not be hitting as intended...
+        # https://github.com/pymmcore-plus/pymmcore-plus/pull/353#issuecomment-2159176491
+        if isinstance(event, SequencedEvent):
+            if event.exposure_sequence:
+                core.stopExposureSequence(self._mmc.getCameraDevice())
+            if event.x_sequence:
+                core.stopXYStageSequence(core.getXYStageDevice())
+            if event.z_sequence:
+                core.stopStageSequence(core.getFocusDevice())
+            for dev, prop in event.property_sequences(core):
+                core.stopPropertySequence(dev, prop)
 
     def teardown_sequence(self, sequence: MDASequence) -> None:
         """Perform any teardown required after the sequence has been executed."""
@@ -346,22 +358,23 @@ class MDAEngine(PMDAEngine):
         cam_device = self._mmc.getCameraDevice()
 
         if event.exposure_sequence:
+            with suppress(RuntimeError):
+                core.stopExposureSequence(cam_device)
             core.loadExposureSequence(cam_device, event.exposure_sequence)
         if event.x_sequence:  # y_sequence is implied and will be the same length
             stage = core.getXYStageDevice()
+            with suppress(RuntimeError):
+                core.stopXYStageSequence(stage)
             core.loadXYStageSequence(stage, event.x_sequence, event.y_sequence)
         if event.z_sequence:
-            # these notes are from Nico Stuurman in AcqEngJ
-            # https://github.com/micro-manager/AcqEngJ/pull/108
-            # at least some zStages freak out (in this case, NIDAQ board) when you
-            # try to load a sequence while the sequence is still running.  Nothing in
-            # the engine stops a stage sequence if all goes well.
-            # Stopping a sequence if it is not running hopefully will not harm anyone.
             zstage = core.getFocusDevice()
-            core.stopStageSequence(zstage)
+            with suppress(RuntimeError):
+                core.stopStageSequence(zstage)
             core.loadStageSequence(zstage, event.z_sequence)
         if prop_seqs := event.property_sequences(core):
             for (dev, prop), value_sequence in prop_seqs.items():
+                with suppress(RuntimeError):
+                    core.stopPropertySequence(dev, prop)
                 core.loadPropertySequence(dev, prop, value_sequence)
 
         # TODO: SLM

{pymmcore_plus-0.9.5 → pymmcore_plus-0.10.1}/src/pymmcore_plus/mda/_runner.py

@@ -89,7 +89,7 @@ class MDARunner:
         if self.is_running():  # pragma: no cover
             raise RuntimeError(
                 "Cannot register a new engine when the current engine is running "
-                "an acquistion. Please cancel the current engine's acquistion "
+                "an acquisition. Please cancel the current engine's acquisition "
                 "before registering"
             )
 
@@ -115,7 +115,7 @@ class MDARunner:
         return self._signals
 
     def is_running(self) -> bool:
-        """Return True if an acquistion is currently underway.
+        """Return True if an acquisition is currently underway.
 
         This will return True at any point between the emission of the
         [`sequenceStarted`][pymmcore_plus.mda.PMDASignaler.sequenceStarted] and
@@ -125,19 +125,19 @@ class MDARunner:
         Returns
         -------
         bool
-            Whether an acquistion is underway.
+            Whether an acquisition is underway.
         """
         return self._running
 
     def is_paused(self) -> bool:
-        """Return True if the acquistion is currently paused.
+        """Return True if the acquisition is currently paused.
 
         Use `toggle_pause` to change the paused state.
 
         Returns
         -------
         bool
-            Whether the current acquistion is paused.
+            Whether the current acquisition is paused.
         """
         return self._paused
 
@@ -145,7 +145,7 @@ class MDARunner:
         """Cancel the currently running acquisition.
 
         This is a no-op if no acquisition is currently running.
-        If an acquisition is running then this will cancel the acquistion and
+        If an acquisition is running then this will cancel the acquisition and
         a sequenceCanceled signal, followed by a sequenceFinished signal will
         be emitted.
         """
@@ -157,7 +157,7 @@ class MDARunner:
 
         To get whether the acquisition is currently paused use the
         [`is_paused`][pymmcore_plus.mda.MDARunner.is_paused] method. This method is a
-        no-op if no acquistion is currently underway.
+        no-op if no acquisition is currently underway.
         """
         if self.is_running():
             self._paused = not self._paused
@@ -169,7 +169,7 @@ class MDARunner:
         *,
         output: SingleOutput | Sequence[SingleOutput] | None = None,
     ) -> None:
-        """Run the multi-dimensional acquistion defined by `sequence`.
+        """Run the multi-dimensional acquisition defined by `sequence`.
 
         Most users should not use this directly as it will block further
         execution. Instead, use the
@@ -278,17 +278,18 @@ class MDARunner:
             logger.info("%s", event)
             engine.setup_event(event)
 
-            output = engine.exec_event(event) or ()  # in case output is None
-
-            for payload in output:
-                img, event, meta = payload
-                if "PerfCounter" in meta:
-                    meta["ElapsedTime-ms"] = (meta["PerfCounter"] - self._t0) * 1000
-                meta["Event"] = event
-                with exceptions_logged():
-                    self._signals.frameReady.emit(img, event, meta)
-
-            teardown_event(event)
+            try:
+                output = engine.exec_event(event) or ()  # in case output is None
+
+                for payload in output:
+                    img, event, meta = payload
+                    if "PerfCounter" in meta:
+                        meta["ElapsedTime-ms"] = (meta["PerfCounter"] - self._t0) * 1000
+                    meta["Event"] = event
+                    with exceptions_logged():
+                        self._signals.frameReady.emit(img, event, meta)
+            finally:
+                teardown_event(event)
 
     def _prepare_to_run(self, sequence: MDASequence) -> PMDAEngine:
         """Set up for the MDA run.

pymmcore_plus-0.10.1/src/pymmcore_plus/mda/handlers/__init__.py

@@ -0,0 +1,11 @@
+from ._img_sequence_writer import ImageSequenceWriter
+from ._ome_tiff_writer import OMETiffWriter
+from ._ome_zarr_writer import OMEZarrWriter
+from ._tensorstore_handler import TensorStoreHandler
+
+__all__ = [
+    "ImageSequenceWriter",
+    "OMEZarrWriter",
+    "OMETiffWriter",
+    "TensorStoreHandler",
+]

{pymmcore_plus-0.9.5 → pymmcore_plus-0.10.1}/src/pymmcore_plus/mda/handlers/_ome_zarr_writer.py

@@ -70,7 +70,7 @@ class OMEZarrWriter(_5DWriterBase["zarr.Array"]):
     │               └── y
     │                   └── x   # chunks will be each XY plane
     ├── ...
-    ├── pn
+    ├── p<n>
     │   ├── .zarray
     │   ├── .zattrs
     │   └── t...

pymmcore_plus-0.10.1/src/pymmcore_plus/mda/handlers/_tensorstore_handler.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import shutil
+import tempfile
+import warnings
+from itertools import product
+from os import PathLike
+from typing import TYPE_CHECKING, Any, cast
+
+from ._util import position_sizes
+
+if TYPE_CHECKING:
+    from typing import Literal, Mapping, Self, Sequence, TypeAlias
+
+    import numpy as np
+    import tensorstore as ts
+    import useq
+
+    TsDriver: TypeAlias = Literal["zarr", "zarr3", "n5", "neuroglancer_precomputed"]
+    EventKey: TypeAlias = frozenset[tuple[str, int]]
+
+# special dimension label used when _nd_storage is False
+FRAME_DIM = "frame"
+
+
+class TensorStoreHandler:
+    """Tensorstore handler for writing MDA sequences.
+
+    This is a performant and shape-agnostic handler for writing MDA sequences to
+    chunked storages like zarr, n5, backed by tensorstore:
+    <https://google.github.io/tensorstore/>
+
+    By default, the handler will store frames in a zarr array, with a shape of
+    (nframes, *frame_shape) and a chunk size of (1, *frame_shape), i.e. each frame
+    is stored in a separate chunk. To customize shape or chunking, override the
+    `get_full_shape`, `get_chunk_layout`, and `get_index_domain` methods (these
+    may change in the future as we learn to use tensorstore better).
+
+    Parameters
+    ----------
+    driver : TsDriver, optional
+        The driver to use for the tensorstore, by default "zarr".  Must be one of
+        "zarr", "zarr3", "n5", or "neuroglancer_precomputed".
+    kvstore : str | dict | None, optional
+        The key-value store to use for the tensorstore, by default "memory://".
+        A dict might look like {'driver': 'file', 'path': '/path/to/dataset.zarr'}
+        see <https://google.github.io/tensorstore/kvstore/index.html#json-KvStore>
+        for all options. If path is provided, the kvstore will be set to file://path
+    path : str | Path | None, optional
+        Convenience for specifying a local filepath. If provided, overrides the
+        kvstore option, to be `file://file_path`.
+    delete_existing : bool, optional
+        Whether to delete the existing dataset if it exists, by default False.
+    spec : Mapping, optional
+        A spec to use when opening the tensorstore, by default None. Values provided
+        in this object will override the default values provided by the handler.
+        This is a complex object that can completely define the tensorstore, see
+        <https://google.github.io/tensorstore/spec.html> for more information.
+
+    Examples
+    --------
+    ```python
+    from pymmcore_plus import CMMCorePlus
+    from pymmcore_plus.mda.handlers import TensorStoreHandler
+    from useq import MDASequence
+
+    core = CMMCorePlus.instance()
+    core.loadSystemConfiguration()
+
+    sequence = MDASequence(
+        channels=["DAPI", {"config": "FITC", "exposure": 1}],
+        stage_positions=[{"x": 1, "y": 1, "name": "some position"}, {"x": 0, "y": 0}],
+        time_plan={"interval": 2, "loops": 3},
+        z_plan={"range": 4, "step": 0.5},
+        axis_order="tpcz",
+    )
+
+    writer = TensorStoreHandler(path="example_ts.zarr", delete_existing=True)
+    core.mda.run(sequence, output=writer)
+    ```
+
+    """
+
+    def __init__(
+        self,
+        *,
+        driver: TsDriver = "zarr",
+        kvstore: str | dict | None = "memory://",
+        path: str | PathLike | None = None,
+        delete_existing: bool = False,
+        spec: Mapping | None = None,
+    ) -> None:
+        try:
+            import tensorstore
+        except ImportError as e:
+            raise ImportError("Tensorstore is required to use this handler.") from e
+
+        self._ts = tensorstore
+
+        self.ts_driver = driver
+        self.kvstore = f"file://{path}" if path is not None else kvstore
+        self.delete_existing = delete_existing
+        self.spec = spec
+
+        # storage of individual frame metadata
+        # maps position key to list of frame metadata
+        self.frame_metadatas: list[tuple[useq.MDAEvent, dict]] = []
+
+        self._size_increment = 300
+
+        self._store: ts.TensorStore | None = None
+        self._futures: list[ts.Future] = []
+        self._frame_indices: dict[EventKey, int | ts.DimExpression] = {}
+
+        # "_nd_storage" means we're greedily attempting to store the data in a
+        # multi-dimensional format based on the axes of the sequence.
+        # for non-deterministic experiments, this often won't work...
+        # _nd_storage False means we simply store data as a 3D array of shape
+        # (nframes, y, x).  `_nd_storage` is set when a new_store is created.
+        self._nd_storage: bool = True
+        self._frame_index: int = 0
+
+        # the highest index seen for each axis
+        self._axis_max: dict[str, int] = {}
+
+    @property
+    def store(self) -> ts.TensorStore | None:
+        """The current tensorstore."""
+        return self._store
+
+    @classmethod
+    def in_tmpdir(
+        cls,
+        suffix: str | None = "",
+        prefix: str | None = "pymmcore_zarr_",
+        dir: str | PathLike[str] | None = None,
+        cleanup_atexit: bool = True,
+        **kwargs: Any,
+    ) -> Self:
+        """Create TensorStoreHandler that writes to a temporary directory.
+
+        Parameters
+        ----------
+        suffix : str, optional
+            If suffix is specified, the file name will end with that suffix, otherwise
+            there will be no suffix.
+        prefix : str, optional
+            If prefix is specified, the file name will begin with that prefix, otherwise
+            a default prefix is used.
+        dir : str or PathLike, optional
+            If dir is specified, the file will be created in that directory, otherwise
+            a default directory is used (tempfile.gettempdir())
+        cleanup_atexit : bool, optional
+            Whether to automatically cleanup the temporary directory when the python
+            process exits. Default is True.
+        **kwargs
+            Remaining kwargs are passed to `TensorStoreHandler.__init__`
+        """
+        # same as zarr.storage.TempStore, but with option not to cleanup
+        path = tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=dir)
+        if cleanup_atexit:
+
+            @atexit.register
+            def _atexit_rmtree(_path: str = path) -> None:  # pragma: no cover
+                if os.path.isdir(_path):
+                    shutil.rmtree(_path, ignore_errors=True)
+
+        return cls(path=path, **kwargs)
+
+    def sequenceStarted(self, seq: useq.MDASequence) -> None:
+        """On sequence started, simply store the sequence."""
+        self._frame_index = 0
+        self._store = None
+        self._futures.clear()
+        self.frame_metadatas.clear()
+        self.current_sequence = seq
+
+    def sequenceFinished(self, seq: useq.MDASequence) -> None:
+        """On sequence finished, clear the current sequence."""
+        if self._store is None:
+            return  # pragma: no cover
+
+        while self._futures:
+            self._futures.pop().result()
+        if not self._nd_storage:
+            self._store = self._store.resize(
+                exclusive_max=(self._frame_index, *self._store.shape[-2:])
+            ).result()
+        if self.frame_metadatas:
+            self.finalize_metadata()
+
+    def frameReady(self, frame: np.ndarray, event: useq.MDAEvent, meta: dict) -> None:
+        """Write frame to the zarr array for the appropriate position."""
+        if self._store is None:
+            self._store = self.new_store(frame, event.sequence, meta).result()
+
+        ts_index: ts.DimExpression | int
+        if self._nd_storage:
+            ts_index = self._event_index_to_store_index(event.index)
+        else:
+            if self._frame_index >= self._store.shape[0]:
+                self._store = self._expand_store(self._store).result()
+            ts_index = self._frame_index
+            # store reverse lookup of event.index -> frame_index
+            self._frame_indices[frozenset(event.index.items())] = ts_index
+
+        # write the new frame asynchronously
+        self._futures.append(self._store[ts_index].write(frame))
+
+        # store, but do not process yet, the frame metadata
+        self.frame_metadatas.append((event, meta))
+        # update the frame counter
+        self._frame_index += 1
+        # remember the highest index seen for each axis
+        for k, v in event.index.items():
+            self._axis_max[k] = max(self._axis_max.get(k, 0), v)
+
+    def isel(
+        self,
+        indexers: Mapping[str, int | slice] | None = None,
+        **indexers_kwargs: int | slice,
+    ) -> np.ndarray:
+        """Select data from the array."""
+        # FIXME: will fail on slices
+        indexers = {**(indexers or {}), **indexers_kwargs}
+        ts_index = self._event_index_to_store_index(indexers)
+        return self._store[ts_index].read().result().squeeze()  # type: ignore
+
+    def new_store(
+        self, frame: np.ndarray, seq: useq.MDASequence | None, meta: dict
+    ) -> ts.Future[ts.TensorStore]:
+        shape, chunks, labels = self.get_shape_chunks_labels(frame.shape, seq)
+        self._nd_storage = FRAME_DIM not in labels
+        return self._ts.open(
+            self.get_spec(),
+            create=True,
+            delete_existing=self.delete_existing,
+            dtype=self._ts.dtype(frame.dtype),
+            shape=shape,
+            chunk_layout=self._ts.ChunkLayout(chunk_shape=chunks),
+            domain=self._ts.IndexDomain(labels=labels),
+        )
+
+    def get_shape_chunks_labels(
+        self, frame_shape: tuple[int, ...], seq: useq.MDASequence | None
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[str, ...]]:
+        labels: tuple[str, ...]
+        if seq is not None and seq.sizes:
+            # expand the sizes to include the largest size we encounter for each axis
+            # in the case of positions with subsequences, we'll still end up with a
+            # jagged array, but it won't take extra space, and we won't get index errors
+            max_sizes = seq.sizes.copy()
+            for psize in position_sizes(seq):
+                for k, v in psize.items():
+                    max_sizes[k] = max(max_sizes.get(k, 0), v)
+
+            # remove axes with length 0
+            labels, sizes = zip(*(x for x in max_sizes.items() if x[1]))
+            full_shape: tuple[int, ...] = (*sizes, *frame_shape)
+        else:
+            labels = (FRAME_DIM,)
+            full_shape = (self._size_increment, *frame_shape)
+
+        chunks = [1] * len(full_shape)
+        chunks[-len(frame_shape) :] = frame_shape
+        labels = (*labels, "y", "x")
+        return full_shape, tuple(chunks), labels
+
+    def get_spec(self) -> dict:
+        """Construct the tensorstore spec."""
+        spec = {"driver": self.ts_driver, "kvstore": self.kvstore}
+        if self.spec:
+            _merge_nested_dicts(spec, self.spec)
+
+        # HACK
+        if self.ts_driver == "zarr":
+            meta = cast(dict, spec.setdefault("metadata", {}))
+            if "dimension_separator" not in meta:
+                meta["dimension_separator"] = "/"
+        return spec
+
+    def finalize_metadata(self) -> None:
+        """Finalize and flush metadata to storage."""
+        if not (store := self._store) or not store.kvstore:
+            return  # pragma: no cover
+
+        data = []
+        for event, meta in self.frame_metadatas:
+            # FIXME: unnecessary ser/des
+            js = event.model_dump_json(exclude={"sequence"}, exclude_defaults=True)
+            meta["Event"] = json.loads(js)
+            data.append(meta)
+
+        metadata = {"frame_metadatas": data}
+        if not self._nd_storage:
+            metadata["frame_indices"] = [
+                (tuple(dict(k).items()), v)  # type: ignore
+                for k, v in self._frame_indices.items()
+            ]
+
+        if self.ts_driver.startswith("zarr"):
+            store.kvstore.write(".zattrs", json.dumps(metadata))
+        elif self.ts_driver == "n5":  # pragma: no cover
+            attrs = json.loads(store.kvstore.read("attributes.json").result().value)
+            attrs.update(metadata)
+            store.kvstore.write("attributes.json", json.dumps(attrs))
+
+    def _expand_store(self, store: ts.TensorStore) -> ts.Future[ts.TensorStore]:
+        """Grow the store by `self._size_increment` frames.
+
+        This is used when _nd_storage mode is False and we've run out of space.
+        """
+        new_shape = [self._frame_index + self._size_increment, *store.shape[-2:]]
+        return store.resize(exclusive_max=new_shape, expand_only=True)
+
+    def _event_index_to_store_index(
+        self, index: Mapping[str, int | slice]
+    ) -> ts.DimExpression:
+        """Convert event index to store index.
+
+        The return value is safe to use as an index to self._store[...]
+        """
+        if self._nd_storage:
+            return self._ts.d[index][tuple(index.values())]
+
+        if any(isinstance(v, slice) for v in index.values()):
+            idx: list | int | ts.DimExpression = self._get_frame_indices(index)
+        else:
+            try:
+                idx = self._frame_indices[frozenset(index.items())]  # type: ignore
+            except KeyError as e:
+                raise KeyError(f"Index {index} not found in frame_indices.") from e
+        return self._ts.d[FRAME_DIM][idx]
+
+    def _get_frame_indices(self, indexers: Mapping[str, int | slice]) -> list[int]:
+        """Convert indexers (with slices) to a list of frame indices."""
+        # converting slice objects to actual indices
+        axis_indices: dict[str, Sequence[int]] = {}
+        for k, v in indexers.items():
+            if isinstance(v, slice):
+                axis_indices[k] = tuple(range(*v.indices(self._axis_max.get(k, 0) + 1)))
+            else:
+                axis_indices[k] = (v,)
+
+        indices: list[int] = []
+        for p in product(*axis_indices.values()):
+            key = frozenset(dict(zip(axis_indices.keys(), p)).items())
+            try:
+                indices.append(self._frame_indices[key])
+            except KeyError:  # pragma: no cover
+                warnings.warn(
+                    f"Index {dict(key)} not found in frame_indices.", stacklevel=2
+                )
+        return indices
+
+
+def _merge_nested_dicts(dict1: dict, dict2: Mapping) -> None:
+    """Merge two nested dictionaries.
+
+    Values in dict2 will override values in dict1.
+    """
+    for key, value in dict2.items():
+        if key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict):
+            _merge_nested_dicts(dict1[key], value)
+        else:
+            dict1[key] = value

pymmcore_plus-0.9.5/tests/io/test_ome_zarr.py → pymmcore_plus-0.10.1/tests/io/test_zarr_writers.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+from queue import Queue
 from typing import TYPE_CHECKING
 
 import numpy as np
 import pytest
 import useq
-from pymmcore_plus.mda.handlers import OMEZarrWriter
+from pymmcore_plus.mda.handlers import OMEZarrWriter, TensorStoreHandler
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -65,18 +66,18 @@ COMPLEX_EXPECTATION = {
 }
 
 
-@pytest.mark.parametrize(
-    "store, mda, expected_shapes",
-    [
-        (None, SIMPLE_MDA, SIMPLE_EXPECTATION),
-        (None, MULTIPOINT_MDA, MULTIPOINT_EXPECTATION),
-        (None, FULL_MDA, FULL_EXPECTATION),
-        ("out.zarr", FULL_MDA, FULL_EXPECTATION),
-        (None, FULL_MDA, FULL_EXPECTATION),
-        ("tmp", FULL_MDA, FULL_EXPECTATION),
-        (None, COMPLEX_MDA, COMPLEX_EXPECTATION),
-    ],
-)
+CASES: list[str | None, useq.MDASequence, dict[str, dict]] = [
+    (None, SIMPLE_MDA, SIMPLE_EXPECTATION),
+    (None, MULTIPOINT_MDA, MULTIPOINT_EXPECTATION),
+    (None, FULL_MDA, FULL_EXPECTATION),
+    ("out.zarr", FULL_MDA, FULL_EXPECTATION),
+    (None, FULL_MDA, FULL_EXPECTATION),
+    ("tmp", FULL_MDA, FULL_EXPECTATION),
+    (None, COMPLEX_MDA, COMPLEX_EXPECTATION),
+]
+
+
+@pytest.mark.parametrize("store, mda, expected_shapes", CASES)
 def test_ome_zarr_writer(
     store: str | None,
     mda: useq.MDASequence,
@@ -128,3 +129,74 @@ def test_ome_zarr_writer(
 
     # smoke test the isel method
     assert isinstance(writer.isel(p=0, t=0, x=slice(0, 100)), np.ndarray)
+
+
+@pytest.mark.parametrize("store, mda, expected_shapes", CASES)
+def test_tensorstore_writer(
+    store: str | None,
+    mda: useq.MDASequence,
+    expected_shapes: dict[str, dict],
+    tmp_path: Path,
+    core: CMMCorePlus,
+) -> None:
+    if store == "tmp":
+        writer = TensorStoreHandler.in_tmpdir()
+    elif store is None:
+        writer = TensorStoreHandler()
+    else:
+        writer = TensorStoreHandler(path=tmp_path / store)
+
+    core.mda.run(mda, output=writer)
+
+    assert writer.store is not None
+
+    expected_sizes = {}
+    for sizes in expected_shapes.values():
+        for dim, size in sizes.items():
+            expected_sizes[dim] = max(sizes.get(dim, 0), size)
+    if len(expected_shapes) > 1:
+        expected_sizes["p"] = len(expected_shapes)
+
+    sizes = dict(zip(writer.store.domain.labels, writer.store.shape))
+    assert sizes == expected_sizes
+
+    if store:
+        # ensure that non-memory stores were written to disk
+        ary = zarr.open(writer.store.kvstore.path)
+        # ensure real data was written
+        assert ary.nchunks_initialized > 0
+        assert ary[0, 0].mean() > (ary.fill_value or 0)
+
+    # smoke test the isel method
+    x = writer.isel(t=0, c=0, x=slice(0, 100))
+    assert isinstance(x, np.ndarray)
+    assert x.shape[-1] == 100
+
+
+def test_tensorstore_writer_spec_override(tmp_path: Path) -> None:
+    writer = TensorStoreHandler(
+        path=tmp_path / "test.zarr",
+        spec={"context": {"cache_pool": {"total_bytes_limit": 10000000}}},
+    )
+
+    assert writer.get_spec()["context"]["cache_pool"]["total_bytes_limit"] == 10000000
+
+
+def test_tensorstore_writer_indeterminate(tmp_path: Path, core: CMMCorePlus) -> None:
+    # FIXME: this test is actually throwing difficult-to-debug exceptions
+    # when driver=='zarr'.  It happens when awaiting the result of self._store.resize()
+    # inside of sequenceFinished.
+    writer = TensorStoreHandler()
+
+    que = Queue()
+    thread = core.run_mda(iter(que.get, None), output=writer)
+    for t in range(2):
+        for z in range(2):
+            que.put(useq.MDAEvent(index={"t": t, "z": z, "c": 0}))
+    que.put(None)
+    thread.join()
+
+    assert writer.isel(t=1, z=1, c=0).shape == (512, 512)
+    assert writer.isel(t=1, z=slice(None), c=0).shape == (2, 512, 512)
+    with pytest.raises(KeyError):
+        writer.isel(t=2, z=2, c=0)

pymmcore_plus-0.9.5/src/pymmcore_plus/mda/handlers/__init__.py

@@ -1,5 +0,0 @@
-from ._img_sequence_writer import ImageSequenceWriter
-from ._ome_tiff_writer import OMETiffWriter
-from ._ome_zarr_writer import OMEZarrWriter
-
-__all__ = ["ImageSequenceWriter", "OMEZarrWriter", "OMETiffWriter"]