Pass kwargs to tensorstore (#332)

* pass kwargs to tensorstore

* style

Author: ziw-liu

iohub/ngff/nodes.py

@@ -383,10 +383,17 @@ def dask_array(self):
     def downscale(self):
         raise NotImplementedError
 
-    def tensorstore(self, concurrency: int | None = None):
-        """Open the zarr array as a TensorStore object.
+    def tensorstore(self, **kwargs):
+        """
+        Open the zarr array as a TensorStore object.
         Needs the optional dependency ``tensorstore``.
 
+        Parameters
+        ----------
+        **kwargs : dict, optional
+            Additional keyword arguments to pass to ``tensorstore.open()``,
+            by default None
+
         Returns
         -------
         TensorStore
@@ -401,15 +408,10 @@ def tensorstore(self, concurrency: int | None = None):
                 "path": str(Path(self.store.root) / self.path.strip("/")),
             },
         }
+        if "read" in kwargs or "write" in kwargs:
+            raise ValueError("Cannot override file mode for the Zarr store.")
         zarr_dataset = ts.open(
-            ts_spec,
-            read=True,
-            write=not self.read_only,
-            context=(
-                ts.Context({"data_copy_concurrency": {"limit": concurrency}})
-                if concurrency
-                else None
-            ),
+            ts_spec, read=True, write=not self.read_only, **kwargs
         ).result()
         return zarr_dataset
 

iohub/ngff/utils.py

@@ -196,10 +196,16 @@ def _save_transformed(
     output_time_indices: int | list[int],
 ) -> None:
     # NOTE: use tensorstore due to zarr-python#3221
+    import tensorstore
+
     with open_ome_zarr(
         output_position_path, layout="fov", mode="r+"
     ) as output_dataset:
-        ts = output_dataset.data.tensorstore(concurrency=4)
+        ts = output_dataset.data.tensorstore(
+            context=tensorstore.Context(
+                {"data_copy_concurrency": {"limit": 4}}
+            )
+        )
     ts.oindex[output_time_indices, output_channel_indices].write(
         transformed
     ).result()

tests/ngff/test_ngff.py

@@ -475,11 +475,19 @@ def test_ome_zarr_to_tensorstore(
     channels_and_random_5d, arr_name, version, concurrency
 ):
     """Test `iohub.ngff.Position.data` to tensorstore"""
+    import tensorstore as ts
+
     channel_names, random_5d = channels_and_random_5d
     with _temp_ome_zarr(
         random_5d, channel_names, arr_name, version=version
     ) as dataset:
-        tstore = dataset[arr_name].tensorstore(concurrency=concurrency)
+        tstore = dataset[arr_name].tensorstore(
+            context=(
+                ts.Context({"data_copy_concurrency": {"limit": concurrency}})
+                if concurrency is not None
+                else None
+            )
+        )
         assert_array_equal(tstore, random_5d)
         zeros = np.zeros_like(random_5d)
         tstore[...].write(zeros).result()
@@ -1288,9 +1296,9 @@ def test_hcs_external_reader(tmp_path):
             fov.create_zeros("0", shape=(1, 2, 3, y_size, x_size), dtype=int)
         n_rows = len(dataset.metadata.rows)
         n_cols = len(dataset.metadata.columns)
-    plate = list(
-        ome_zarr.reader.Reader(ome_zarr.io.parse_url(store_path))()
-    )[0]
+    plate = list(ome_zarr.reader.Reader(ome_zarr.io.parse_url(store_path))())[
+        0
+    ]
     assert plate.data[0].shape == (1, 2, 3, y_size * n_rows, x_size * n_cols)
     assert plate.data[0].dtype == int
     assert not plate.data[0].any()