deduplicate storage option util; use chunks from data when not specified in storage options

Author: LucaMarconato

src/spatialdata/_io/io_raster.py

@@ -90,34 +90,27 @@ def _normalize_explicit_chunks(chunks: object) -> tuple[int, ...] | int:
     return normalized
 
 
-def _prepare_single_scale_storage_options(
+def _prepare_storage_options(
     storage_options: JSONDict | list[JSONDict] | None,
-) -> JSONDict | list[JSONDict] | None:
+    data: list[da.Array],
+) -> list[JSONDict]:
     if storage_options is None:
-        return None
-    if isinstance(storage_options, dict):
-        prepared = dict(storage_options)
-        if "chunks" in prepared:
-            prepared["chunks"] = _normalize_explicit_chunks(prepared["chunks"])
-        return prepared
-    return [dict(options) for options in storage_options]
-
-
-def _prepare_multiscale_storage_options(
-    storage_options: JSONDict | list[JSONDict] | None,
-) -> JSONDict | list[JSONDict] | None:
-    if storage_options is None:
-        return None
+        return [{"chunks": _normalize_explicit_chunks(arr.chunks)} for arr in data]
     if isinstance(storage_options, dict):
+        if "chunks" not in storage_options:
+            return [{**storage_options, "chunks": _normalize_explicit_chunks(arr.chunks)} for arr in data]
         prepared = dict(storage_options)
-        if "chunks" in prepared:
-            prepared["chunks"] = _normalize_explicit_chunks(prepared["chunks"])
-        return prepared
-
-    prepared_options = [dict(options) for options in storage_options]
-    for options in prepared_options:
-        if "chunks" in options:
-            options["chunks"] = _normalize_explicit_chunks(options["chunks"])
+        prepared["chunks"] = _normalize_explicit_chunks(prepared["chunks"])
+        return prepared  # type: ignore[return-value]
+
+    prepared_options = []
+    for i, options in enumerate(storage_options):
+        opts = dict(options)
+        if "chunks" not in opts:
+            opts["chunks"] = _normalize_explicit_chunks(data[i].chunks)
+        else:
+            opts["chunks"] = _normalize_explicit_chunks(opts["chunks"])
+        prepared_options.append(opts)
     return prepared_options
 
 
@@ -335,7 +328,7 @@ def _write_raster_dataarray(
         raise ValueError(f"{element_name} does not have any transformations and can therefore not be written.")
     input_axes: tuple[str, ...] = tuple(raster_data.dims)
     parsed_axes = _get_valid_axes(axes=list(input_axes), fmt=raster_format)
-    storage_options = _prepare_single_scale_storage_options(storage_options)
+    storage_options = _prepare_storage_options(storage_options, [data])
     # Explicitly disable pyramid generation for single-scale rasters. Recent ome-zarr versions default
     # write_image()/write_labels() to scale_factors=(2, 4, 8, 16), which would otherwise write s0, s1, ...
     # even when the input is a plain DataArray.
@@ -405,7 +398,7 @@ def _write_raster_datatree(
         raise ValueError(f"{element_name} does not have any transformations and can therefore not be written.")
 
     parsed_axes = _get_valid_axes(axes=list(input_axes), fmt=raster_format)
-    storage_options = _prepare_multiscale_storage_options(storage_options)
+    storage_options = _prepare_storage_options(storage_options, data)
     ome_zarr_format = get_ome_zarr_format(raster_format)
     dask_delayed = write_multi_scale_ngff(
         pyramid=data,

tests/io/test_readwrite.py

@@ -631,7 +631,11 @@ def test_write_irregular_dask_chunks_without_explicit_storage_options(tmp_path:
     image = Image2DModel.parse(data, dims=("c", "y", "x"))
     sdata = SpatialData(images={"image": image})
 
-    sdata.write(tmp_path / "data.zarr")
+    with pytest.raises(
+        ValueError,
+        match="storage_options\\['chunks'\\] must be a Zarr chunk shape or a regular Dask chunk grid",
+    ):
+        sdata.write(tmp_path / "data.zarr")
 
 
 def test_write_image_normalizes_explicit_regular_dask_chunk_grid(tmp_path: Path) -> None: