code edits for zarr v3; tests fail

Author: LucaMarconato

src/spatialdata/_io/io_raster.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 import dask.array as da
 import numpy as np
@@ -231,7 +231,9 @@ def _write_raster(
 
 
 def _apply_compression(
-    storage_options: JSONDict | list[JSONDict], compressor: dict[Literal["lz4", "zstd"], int] | None
+    storage_options: JSONDict | list[JSONDict],
+    compressor: dict[Literal["lz4", "zstd"], int] | None,
+    zarr_format: Literal[2, 3] = 3,
 ) -> JSONDict | list[JSONDict]:
     """Apply compression settings to storage options.
 
@@ -241,23 +243,39 @@ def _apply_compression(
         Storage options for zarr arrays
     compressor
         Compression settings as a dictionary with a single key-value pair
+    zarr_format
+        The zarr format version (2 or 3)
 
     Returns
     -------
     Updated storage options with compression settings
     """
-    from zarr.codecs import Blosc
+    # For zarr disk format v2, use numcodecs.Blosc
+    # For zarr disk format v3, use zarr.codecs.Blosc
+    from numcodecs import Blosc as BloscV2
+    from zarr.codecs import Blosc as BloscV3
 
     if not compressor:
         return storage_options
 
     ((compression, compression_level),) = compressor.items()
 
+    assert BloscV2.SHUFFLE == 1
+    blosc_v2 = BloscV2(cname=compression, clevel=compression_level, shuffle=1)
+    blosc_v3 = BloscV3(cname=compression, clevel=compression_level, shuffle=1)
+
+    def _update_dict(d: dict[str, Any]) -> None:
+        if zarr_format == 2:
+            d["compressor"] = blosc_v2
+        elif zarr_format == 3:
+            d["zarr_array_kwargs"] = {"compressors": [blosc_v3]}
+
     if isinstance(storage_options, dict):
-        storage_options["compressor"] = Blosc(cname=compression, clevel=compression_level, shuffle=1)
-    elif isinstance(storage_options, list):
+        _update_dict(d=storage_options)
+    else:
+        assert isinstance(storage_options, list)
         for option in storage_options:
-            option["compressor"] = Blosc(cname=compression, clevel=compression_level, shuffle=1)
+            _update_dict(d=option)
 
     return storage_options
 
@@ -311,7 +329,9 @@ def _write_raster_dataarray(
         storage_options = {"chunks": chunks}
 
     # Apply compression if specified
-    storage_options = _apply_compression(storage_options, compressor)
+    storage_options = _apply_compression(
+        storage_options, compressor, zarr_format=cast(Literal[2, 3], raster_format.zarr_format)
+    )
 
     # Scaler needs to be None since we are passing the data already downscaled for the multiscale case.
     # We need  this because the argument of write_image_ngff is called image while the argument of
@@ -388,7 +408,7 @@ def _write_raster_datatree(
     else:
         storage_options = [{"chunks": chunk} for chunk in chunks]
     # Apply compression if specified
-    storage_options = _apply_compression(storage_options, compressor)
+    storage_options = _apply_compression(storage_options, compressor, zarr_format=raster_format.zarr_format)
 
     parsed_axes = _get_valid_axes(axes=list(input_axes), fmt=raster_format)
     ome_zarr_format = get_ome_zarr_format(raster_format)

tests/io/test_readwrite.py

@@ -149,6 +149,20 @@ def test_multiple_tables(
         sdata_tables = SpatialData(tables={str(i): tables[i] for i in range(len(tables))})
         self._test_table(tmp_path, sdata_tables, sdata_container_format=sdata_container_format)
 
+    def test_roundtrip(
+        self,
+        tmp_path: str,
+        sdata: SpatialData,
+        sdata_container_format: SpatialDataContainerFormatType,
+    ) -> None:
+        tmpdir = Path(tmp_path) / "tmp.zarr"
+
+        sdata.write(tmpdir, sdata_formats=sdata_container_format)
+        sdata2 = SpatialData.read(tmpdir)
+        tmpdir2 = Path(tmp_path) / "tmp2.zarr"
+        sdata2.write(tmpdir2, sdata_formats=sdata_container_format)
+        _are_directories_identical(tmpdir, tmpdir2, exclude_regexp="[1-9][0-9]*.*")
+
     def test_compression_roundtrip(
         self,
         tmp_path: str,
@@ -168,31 +182,55 @@ def test_compression_roundtrip(
         full_sdata.write(tmpdir, compressor={"zstd": 8}, sdata_formats=sdata_container_format)
 
         # sourcery skip: no-loop-in-tests
-        for element in ["image2d", "image2d_multiscale"]:
-            arr = zarr.open_group(tmpdir / "images", mode="r")[element]["0"]
+        for element in ["image2d", "image2d_multiscale", "labels2d", "labels2d_multiscale"]:
+            element_type = "images" if element.startswith("image") else "labels"
+            arr = zarr.open_group(tmpdir / element_type, mode="r")[element]["0"]
             compressor = arr.compressors[0]
-            assert compressor.cnam == "zstd"
-            assert compressor.clevel == 8
 
-        for element in ["labels2d", "labels2d_multiscale"]:
-            arr = zarr.open_group(tmpdir / "labels", mode="r")[element]["0"]
-            compressor = arr.compressors[0]
-            assert compressor.cname == "zstd"
-            assert compressor.clevel == 8
+            # TODO: all these tests fail because the compression arguments are not passed to Dask
+            if sdata_container_format.zarr_format == 2:
+                assert compressor.cname == "zstd"
+                assert compressor.clevel == 8
+            elif sdata_container_format.zarr_format == 3:
+                from zarr.codecs.zstd import ZstdCodec
 
-    def test_roundtrip(
+                assert isinstance(compressor, ZstdCodec)
+                assert compressor.level == 8
+
+    @pytest.mark.parametrize("compressor", [{"lz4": 3}, {"zstd": 7}])
+    @pytest.mark.parametrize("element", [("images", "image2d"), ("labels", "labels2d")])
+    def test_write_element_compression(
         self,
         tmp_path: str,
-        sdata: SpatialData,
+        full_sdata: SpatialData,
+        compressor: dict[Literal["lz4", "zstd"], int],
+        element: str,
         sdata_container_format: SpatialDataContainerFormatType,
-    ) -> None:
-        tmpdir = Path(tmp_path) / "tmp.zarr"
-
+    ):
+        tmpdir = Path(tmp_path) / "compression.zarr"
+        sdata = SpatialData()
         sdata.write(tmpdir, sdata_formats=sdata_container_format)
-        sdata2 = SpatialData.read(tmpdir)
-        tmpdir2 = Path(tmp_path) / "tmp2.zarr"
-        sdata2.write(tmpdir2, sdata_formats=sdata_container_format)
-        _are_directories_identical(tmpdir, tmpdir2, exclude_regexp="[1-9][0-9]*.*")
+
+        sdata["element"] = full_sdata[element[1]]
+        sdata.write_element("element", compressor=compressor, sdata_formats=sdata_container_format)
+
+        arr = zarr.open_group(tmpdir / element[0], mode="r")["element"]["0"]
+        compression = arr.compressors[0]
+
+        # TODO: all these tests fail because the compression arguments are not passed to Dask
+        if sdata_container_format.zarr_format == 2:
+            assert compression.cname == list(compressor.keys())[0]
+            assert compression.clevel == list(compressor.values())[0]
+        elif sdata_container_format.zarr_format == 3:
+            from zarr.codecs import ZstdCodec
+
+            compressor_name = list(compressor.keys())[0]
+            if compressor_name == "zstd":
+                assert isinstance(compression, ZstdCodec)
+            # TODO: fix
+            # elif compressor_name == 'lz4':
+            #     assert isinstance(compression, ???)
+            assert compression.level == list(compressor.values())[0]
 
     def test_incremental_io_list_of_elements(
         self,
@@ -381,28 +419,6 @@ def test_io_and_lazy_loading_raster(self, images, labels, sdata_container_format
                 assert any("from-zarr" in key for key in dask1.dask.layers)
                 assert len(get_dask_backing_files(sdata2)) > 0
 
-    @pytest.mark.parametrize("compressor", [{"lz4": 3}, {"zstd": 7}])
-    @pytest.mark.parametrize("element", [("images", "image2d"), ("labels", "labels2d")])
-    def test_write_element_compression(
-        self,
-        tmp_path: str,
-        full_sdata: SpatialData,
-        compressor: dict[Literal["lz4", "zstd"], int],
-        element: str,
-        sdata_container_format: SpatialDataContainerFormatType,
-    ):
-        tmpdir = Path(tmp_path) / "compression.zarr"
-        sdata = SpatialData()
-        sdata.write(tmpdir, sdata_formats=sdata_container_format)
-
-        sdata["element"] = full_sdata[element[1]]
-        sdata.write_element("element", compressor=compressor, sdata_formats=sdata_container_format)
-
-        arr = zarr.open_group(tmpdir / element[0], mode="r")["element"]["0"]
-        compression = arr.compressors[0]
-        assert compression.cname == list(compressor.keys())[0]
-        assert compression.clevel == list(compressor.values())[0]
-
     def test_replace_transformation_on_disk_raster(
         self, images, labels, sdata_container_format: SpatialDataContainerFormatType
     ):