refactor: pack zarr-related config into a dataclasss for reuse

linc_convert/modalities/__init__.py

@@ -1,5 +1,4 @@
 """Converters for all imaging modalities."""
 
 __all__ = ["df", "lsm", "wk", "psoct"]
-from . import df, lsm, wk, psoct
-
+from . import df, lsm, psoct, wk

linc_convert/modalities/df/multi_slice.py

@@ -23,6 +23,7 @@
 from linc_convert.utils.math import ceildiv, floordiv
 from linc_convert.utils.orientation import center_affine, orientation_to_affine
 from linc_convert.utils.zarr.compressor import make_compressor
+from linc_convert.utils.zarr.zarr_config import ZarrConfig
 
 HOME = "/space/aspasia/2/users/linc/000003"
 
@@ -38,13 +39,9 @@
 @ms.default
 def convert(
     inp: list[str],
-    out: str | None = None,
     *,
-    chunk: int = 1024,
-    compressor: str = "blosc",
-    compressor_opt: str = "{}",
+    zarr_config: ZarrConfig,
     max_load: int = 16384,
-    nii: bool = False,
     orientation: str = "coronal",
     center: bool = True,
     thickness: float | None = None,
@@ -84,25 +81,21 @@ def convert(
     ----------
     inp
         Path to the input slices
-    out
-        Path to the output Zarr directory [<dirname(INP)>.ome.zarr]
-    chunk
-        Output chunk size
-    compressor : {blosc, zlib, raw}
-        Compression method
-    compressor_opt
-        Compression options
     max_load
         Maximum input chunk size
-    nii
-        Convert to nifti-zarr. True if path ends in ".nii.zarr"
     orientation
         Orientation of the slice
     center
         Set RAS[0, 0, 0] at FOV center
     thickness
         Slice thickness
     """
+    out: str = zarr_config.out
+    chunk: int = zarr_config.chunk[0]
+    compressor: str = zarr_config.compressor
+    compressor_opt: str = zarr_config.compressor_opt
+    nii: bool = zarr_config.nii
+
     # Default output path
     if not out:
         out = os.path.splitext(inp[0])[0]

linc_convert/modalities/df/single_slice.py

@@ -22,6 +22,7 @@
 from linc_convert.utils.math import ceildiv
 from linc_convert.utils.orientation import center_affine, orientation_to_affine
 from linc_convert.utils.zarr.compressor import make_compressor
+from linc_convert.utils.zarr.zarr_config import ZarrConfig
 
 ss = App(name="singleslice", help_format="markdown")
 df.command(ss)
@@ -30,13 +31,9 @@
 @ss.default
 def convert(
     inp: str,
-    out: str | None = None,
     *,
-    chunk: int = 1024,
-    compressor: str = "blosc",
-    compressor_opt: str = "{}",
+    zarr_config: ZarrConfig,
     max_load: int = 16384,
-    nii: bool = False,
     orientation: str = "coronal",
     center: bool = True,
     thickness: float | None = None,
@@ -71,25 +68,21 @@ def convert(
     ----------
     inp
         Path to the input JP2 file
-    out
-        Path to the output Zarr directory [<INP>.ome.zarr]
-    chunk
-        Output chunk size
-    compressor : {blosc, zlib, raw}
-        Compression method
-    compressor_opt
-        Compression options
     max_load
         Maximum input chunk size
-    nii
-        Convert to nifti-zarr. True if path ends in ".nii.zarr"
     orientation
         Orientation of the slice
     center
         Set RAS[0, 0, 0] at FOV center
     thickness
         Slice thickness
     """
+    out: str = zarr_config.out
+    chunk: int = zarr_config.chunk
+    compressor: str = zarr_config.compressor
+    compressor_opt: str = zarr_config.compressor_opt
+    nii: bool = zarr_config.nii
+
     if not out:
         out = os.path.splitext(inp)[0]
         out += ".nii.zarr" if nii else ".ome.zarr"

linc_convert/modalities/lsm/mosaic.py

@@ -23,6 +23,7 @@
 from linc_convert.utils.math import ceildiv
 from linc_convert.utils.orientation import center_affine, orientation_to_affine
 from linc_convert.utils.zarr.compressor import make_compressor
+from linc_convert.utils.zarr.zarr_config import ZarrConfig
 
 mosaic = cyclopts.App(name="mosaic", help_format="markdown")
 lsm.command(mosaic)
@@ -31,13 +32,9 @@
 @mosaic.default
 def convert(
     inp: str,
-    out: str = None,
     *,
-    chunk: int = 128,
-    compressor: str = "blosc",
-    compressor_opt: str = "{}",
+    zarr_config: ZarrConfig,
     max_load: int = 512,
-    nii: bool = False,
     orientation: str = "coronal",
     center: bool = True,
     thickness: float | None = None,
@@ -72,25 +69,21 @@ def convert(
         Path to the root directory, which contains a collection of
         subfolders named `*_z{:02d}_y{:02d}*`, each containing a
         collection of files named `*_plane{:03d}_c{:d}.tiff`.
-    out
-        Path to the output Zarr directory [<INP>.ome.zarr]
-    chunk
-        Output chunk size
-    compressor : {blosc, zlib, raw}
-        Compression method
-    compressor_opt
-        Compression options
     max_load
         Maximum input chunk size when building pyramid
-    nii
-        Convert to nifti-zarr. True if path ends in ".nii.zarr".
     orientation
         Orientation of the slice
     center
         Set RAS[0, 0, 0] at FOV center
     voxel_size
         Voxel size along the X, Y and Z dimension, in micron.
     """
+    out: str = zarr_config.out
+    chunk: int = zarr_config.chunk[0]
+    compressor: str = zarr_config.compressor
+    compressor_opt: str = zarr_config.compressor_opt
+    nii: bool = zarr_config.nii
+
     if isinstance(compressor_opt, str):
         compressor_opt = ast.literal_eval(compressor_opt)
 

linc_convert/modalities/psoct/multi_slice.py

@@ -31,6 +31,7 @@
 from linc_convert.utils.orientation import center_affine, orientation_to_affine
 from linc_convert.utils.unit import to_nifti_unit, to_ome_unit
 from linc_convert.utils.zarr.compressor import make_compressor
+from linc_convert.utils.zarr.zarr_config import ZarrConfig, _ZarrConfig
 
 multi_slice = cyclopts.App(name="multi_slice", help_format="markdown")
 psoct.command(multi_slice)
@@ -40,13 +41,12 @@ def _automap(func: Callable) -> Callable:
     """Automatically maps the array in the mat file."""
 
     @wraps(func)
-    def wrapper(inp: list[str], out: str = None, **kwargs: dict) -> callable:
-        if out is None:
-            out = os.path.splitext(inp[0])[0]
-            out += ".nii.zarr" if kwargs.get("nii", False) else ".ome.zarr"
-        kwargs["nii"] = kwargs.get("nii", False) or out.endswith(".nii.zarr")
+    def wrapper(inp: list[str], zarr_config: _ZarrConfig, **kwargs: dict) -> callable:
+        if zarr_config.out is None:
+            zarr_config.out = os.path.splitext(inp[0])[0]
+            zarr_config.out += ".nii.zarr" if zarr_config.nii else ".ome.zarr"
         dat = _mapmat(inp, kwargs.get("key", None))
-        return func(dat, out, **kwargs)
+        return func(dat, zarr_config=zarr_config, **kwargs)
 
     return wrapper
 
@@ -163,17 +163,13 @@ def make_wrapper(fname: str) -> callable:
 @_automap
 def convert(
     inp: list[str],
-    out: Optional[str] = None,
     *,
+    zarr_config: ZarrConfig,
     key: Optional[str] = None,
     meta: str = None,
-    chunk: int = 128,
-    compressor: str = "blosc",
-    compressor_opt: str = "{}",
     max_load: int = 128,
     max_levels: int = 5,
     no_pool: Optional[int] = None,
-    nii: bool = False,
     orientation: str = "RAS",
     center: bool = True,
     dtype: str | None = None,
@@ -192,33 +188,29 @@ def convert(
     ----------
     inp
         Path to the input mat file
-    out
-        Path to the output Zarr directory [<INP>.ome.zarr]
     key
         Key of the array to be extracted, default to first key found
     meta
         Path to the metadata file
-    chunk
-        Output chunk size
-    compressor : {blosc, zlib, raw}
-        Compression method
-    compressor_opt
-        Compression options
     max_load
         Maximum input chunk size
     max_levels
         Maximum number of pyramid levels
     no_pool
         Index of dimension to not pool when building pyramid.
-    nii
-        Convert to nifti-zarr. True if path ends in ".nii.zarr"
     orientation
         Orientation of the volume
     center
         Set RAS[0, 0, 0] at FOV center
     dtype
         Data type to write into
     """
+    out: str = zarr_config.out
+    chunk: int = zarr_config.chunk[0]
+    compressor: str = zarr_config.compressor
+    compressor_opt: str = zarr_config.compressor_opt
+    nii: bool = zarr_config.nii
+
     if isinstance(compressor_opt, str):
         compressor_opt = ast.literal_eval(compressor_opt)
 

linc_convert/modalities/psoct/single_volume.py

@@ -32,6 +32,7 @@
 from linc_convert.utils.orientation import center_affine, orientation_to_affine
 from linc_convert.utils.unit import to_nifti_unit, to_ome_unit
 from linc_convert.utils.zarr.compressor import make_compressor
+from linc_convert.utils.zarr.zarr_config import ZarrConfig, _ZarrConfig
 
 single_volume = cyclopts.App(name="single_volume", help_format="markdown")
 psoct.command(single_volume)
@@ -41,13 +42,13 @@ def _automap(func: Callable) -> Callable:
     """Automatically map the array in the mat file."""
 
     @wraps(func)
-    def wrapper(inp: str, out: str = None, **kwargs: dict) -> None:
-        if out is None:
-            out = os.path.splitext(inp[0])[0]
-            out += ".nii.zarr" if kwargs.get("nii", False) else ".ome.zarr"
-        kwargs["nii"] = kwargs.get("nii", False) or out.endswith(".nii.zarr")
+    def wrapper(inp: str, zarr_config: _ZarrConfig, **kwargs: dict) -> None:
+        if zarr_config.out is None:
+            zarr_config.out = os.path.splitext(inp[0])[0]
+            zarr_config.out += ".nii.zarr" if zarr_config.nii else ".ome.zarr"
+        # kwargs["nii"] = kwargs.get("nii", False) or out.endswith(".nii.zarr")
         with _mapmat(inp, kwargs.get("key", None)) as dat:
-            return func(dat, out, **kwargs)
+            return func(dat, zarr_config=zarr_config, **kwargs)
 
     return wrapper
 
@@ -86,17 +87,13 @@ def _mapmat(fname: str, key: str = None) -> None:
 @_automap
 def convert(
     inp: str,
-    out: Optional[str] = None,
     *,
+    zarr_config: ZarrConfig,
     key: Optional[str] = None,
     meta: str = None,
-    chunk: int = 128,
-    compressor: str = "blosc",
-    compressor_opt: str = "{}",
     max_load: int = 128,
     max_levels: int = 5,
     no_pool: Optional[int] = None,
-    nii: bool = False,
     orientation: str = "RAS",
     center: bool = True,
 ) -> None:
@@ -135,6 +132,12 @@ def convert(
     center
         Set RAS[0, 0, 0] at FOV center
     """
+    out: str = zarr_config.out
+    chunk: int = zarr_config.chunk
+    compressor: str = zarr_config.compressor
+    compressor_opt: str = zarr_config.compressor_opt
+    nii: bool = zarr_config.nii
+
     if isinstance(compressor_opt, str):
         compressor_opt = ast.literal_eval(compressor_opt)
 

linc_convert/utils/zarr/zarr_config.py

@@ -1,7 +1,7 @@
 """Configuration related to output Zarr Archive."""
 
 from dataclasses import dataclass
-from typing import Annotated, Literal
+from typing import Annotated, Literal, Optional
 
 from cyclopts import Parameter
 
@@ -38,7 +38,7 @@ class _ZarrConfig:
 
     """
 
-    out: str = ""
+    out: Optional[str] = None
     chunk: tuple[int] = (128,)
     shard: list[int | str] | None = None
     version: Literal[2, 3] = 3

tests/test_df.py

@@ -7,6 +7,7 @@
 from helper import _cmp_zarr_archives
 
 from linc_convert.modalities.df import multi_slice
+from linc_convert.utils.zarr.zarr_config import _ZarrConfig
 
 
 def _write_test_data(directory: str) -> None:
@@ -25,5 +26,5 @@ def test_df(tmp_path):
     output_zarr = tmp_path / "output.zarr"
     files = glob.glob(os.path.join(tmp_path, "*.jp2"))
     files.sort()
-    multi_slice.convert(files, str(output_zarr))
+    multi_slice.convert(files, zarr_config=_ZarrConfig(str(output_zarr)))
     assert _cmp_zarr_archives(str(output_zarr), "data/df.zarr.zip")

tests/test_lsm.py

@@ -5,6 +5,7 @@
 from helper import _cmp_zarr_archives
 
 from linc_convert.modalities.lsm import mosaic
+from linc_convert.utils.zarr.zarr_config import _ZarrConfig
 
 
 def _write_test_data(directory: str) -> None:
@@ -24,5 +25,5 @@ def _write_test_data(directory: str) -> None:
 def test_lsm(tmp_path):
     _write_test_data(tmp_path)
     output_zarr = tmp_path / "output.zarr"
-    mosaic.convert(str(tmp_path), str(output_zarr))
+    mosaic.convert(str(tmp_path), zarr_config=_ZarrConfig(str(output_zarr)))
     assert _cmp_zarr_archives(str(output_zarr), "data/lsm.zarr.zip")