fix: compressor=None kwarg handling with zarr v3 (#2270)

Author: ilan-gold

docs/release-notes/2270.fix.md

@@ -0,0 +1 @@
+Fix `compressor` kwarg handling when writing to zarr v3 {user}`ilan-gold`

src/anndata/_io/specs/methods.py

@@ -94,9 +94,21 @@
 #     return False
 
 
-def zarr_v3_compressor_compat(dataset_kwargs) -> dict:
-    if not is_zarr_v2() and (compressor := dataset_kwargs.pop("compressor", None)):
-        dataset_kwargs["compressors"] = compressor
+def zarr_v3_compressor_compat(dataset_kwargs: dict) -> dict:
+    """Handle mismatch between our compressor kwarg and :func:`zarr.create_array` in v3's `compressors` arg
+    See https://zarr.readthedocs.io/en/stable/api/zarr/create/#zarr.create_array
+
+    Parameters
+    ----------
+    dataset_kwarg
+        The kwarg dict potentially containing "compressor"
+
+    Returns
+    -------
+        The kwarg dict with "compressor" moved to "compressors" if zarr v3 is in use.
+    """
+    if not is_zarr_v2() and "compressor" in dataset_kwargs:
+        dataset_kwargs["compressors"] = dataset_kwargs.pop("compressor")
     return dataset_kwargs
 
 

tests/test_readwrite.py

@@ -377,13 +377,20 @@ def check_compressed(key, value):
 
 
 @pytest.mark.parametrize("zarr_write_format", [2, 3])
-def test_zarr_compression(tmp_path, zarr_write_format):
+@pytest.mark.parametrize(
+    "use_compression", [True, False], ids=["compressed", "uncompressed"]
+)
+def test_zarr_compression(
+    tmp_path: Path, zarr_write_format: Literal[2, 3], *, use_compression: bool
+):
     if zarr_write_format == 3 and is_zarr_v2():
         pytest.xfail("Cannot write zarr v3 format with v2 package")
     ad.settings.zarr_write_format = zarr_write_format
     pth = str(Path(tmp_path) / "adata.zarr")
     adata = gen_adata((10, 8), **GEN_ADATA_NO_XARRAY_ARGS)
-    if zarr_write_format == 2 or is_zarr_v2():
+    if not use_compression:
+        compressor = None
+    elif zarr_write_format == 2 or is_zarr_v2():
         from numcodecs import Blosc
 
         compressor = Blosc(cname="zstd", clevel=3, shuffle=Blosc.BITSHUFFLE)
@@ -393,21 +400,23 @@ def test_zarr_compression(tmp_path, zarr_write_format):
         # Don't use Blosc since it's defaults can change:
         # https://github.com/zarr-developers/zarr-python/pull/3545
         compressor = ZstdCodec(level=3, checksum=True)
-    not_compressed = []
+    wrongly_compressed = []
 
     ad.io.write_zarr(pth, adata, compressor=compressor)
 
     def check_compressed(value, key):
         if not isinstance(value, ZarrArray) or value.shape == ():
             return None
-        (read_compressor,) = value.compressors
+        (read_compressor,) = value.compressors or [None]
         if zarr_write_format == 2:
             if read_compressor != compressor:
-                not_compressed.append(key)
+                wrongly_compressed.append(key)
             return None
-        if read_compressor.to_dict() != compressor.to_dict():
-            print(read_compressor.to_dict(), compressor.to_dict())
-            not_compressed.append(key)
+        if (compressor is None and read_compressor is not None) or (
+            None not in {compressor, read_compressor}
+            and read_compressor.to_dict() != compressor.to_dict()
+        ):
+            wrongly_compressed.append(key)
 
     if is_zarr_v2():
         with zarr.open(pth, "r") as f:
@@ -416,14 +425,7 @@ def check_compressed(value, key):
         f = zarr.open(pth, mode="r")
         for key, value in f.members(max_depth=None):
             check_compressed(value, key)
-
-    if not_compressed:
-        sep = "\n\t"
-        msg = (
-            f"These elements were not compressed correctly:{sep}"
-            f"{sep.join(not_compressed)}"
-        )
-        raise AssertionError(msg)
+    assert not wrongly_compressed, "Some elements were not (un)compressed correctly"
 
     expected = ad.read_zarr(pth)
     assert_equal(adata, expected)