Merge pull request #365 from dstansby/isotropic-chunks

Fix scalar chunking

CHANGELOG.md

@@ -3,6 +3,11 @@
 - Correctly specify maximum compatible fsspec version. ([#338](https://github.com/ome/ome-zarr-py/pull/338))
 - Add tests on Python 3.12. ([#338](https://github.com/ome/ome-zarr-py/pull/338))
 - Write OMERO metadata. ([#261](https://github.com/ome/ome-zarr-py/pull/261))
+- Fixed chunking when a scalar value for chunks is given. Previously
+  passing ``storage_options={"chunks": <int>}`` only set the chunk
+  size of the final dimension. Now the chunk size of all dimensions is
+  set to this value, which is identical behaviour to ``zarr-python``.
+  ([#365](https://github.com/ome/ome-zarr-py/pull/365))
 
 # 0.8.3 (November 2023)
 

ome_zarr/writer.py

@@ -921,14 +921,13 @@ def _retuple(
 
     E.g. if chunks is (64, 64) and shape is (3, 4, 5, 1028, 1028)
     return (3, 4, 5, 64, 64)
+
+    If chunks is an integer, it is applied to all dimensions, to match
+    the behaviour of zarr-python.
     """
 
-    _chunks: Tuple[Any, ...]
     if isinstance(chunks, int):
-        _chunks = (chunks,)
-    else:
-        _chunks = chunks
-
-    dims_to_add = len(shape) - len(_chunks)
+        return tuple([chunks] * len(shape))
 
-    return (*shape[:dims_to_add], *_chunks)
+    dims_to_add = len(shape) - len(chunks)
+    return (*shape[:dims_to_add], *chunks)

tests/test_writer.py

@@ -216,6 +216,20 @@ def test_write_image_dask(self, read_from_zarr, compute):
                 shallow=False,
             )
 
+    def test_write_image_scalar_chunks(self):
+        """
+        Make sure a scalar chunks value is applied to all dimensions,
+        matching the behaviour of zarr-python.
+        """
+        shape = (64, 64, 64)
+        data = np.array(self.create_data(shape))
+        write_image(
+            image=data, group=self.group, axes="xyz", storage_options={"chunks": 32}
+        )
+        for data in self.group.values():
+            print(data)
+            assert data.chunks == (32, 32, 32)
+
     @pytest.mark.parametrize("array_constructor", [np.array, da.from_array])
     def test_write_image_compressed(self, array_constructor):
         shape = (64, 64, 64)