Handle overwrite for existing pyramid levels & add test

Author: jluethi

fractal_tasks_core/pyramids.py

@@ -22,7 +22,7 @@
 
 import dask.array as da
 import numpy as np
-import zarr
+import zarr.errors
 
 logger = logging.getLogger(__name__)
 
@@ -107,6 +107,18 @@ def build_pyramid(
         if open_array_kwargs is None:
             open_array_kwargs = {}
 
+        # If overwrite is false, check that the array doesn't exist yet
+        if not overwrite:
+            try:
+                zarr.open(f"{zarrurl}/{ind_level}", mode="r")
+                raise ValueError(
+                    f"While building the pyramids, pyramid level {ind_level} "
+                    "already existed, but `build_pyramid` was called with "
+                    f"{overwrite=}."
+                )
+            except zarr.errors.PathNotFoundError:
+                pass
+
         zarrarr = zarr.open(
             f"{zarrurl}/{ind_level}",
             shape=newlevel_rechunked.shape,

tests/test_unit_pyramid_creation.py

@@ -81,3 +81,19 @@ def test_build_pyramid(tmp_path):
     # check that the empty chunks are not written to disk
     assert not (tmp_path / "F.zarr/1/0.0").exists()
     assert not (tmp_path / "F.zarr/2/0.0").exists()
+
+
+def test_build_pyramid_overwrite(tmp_path):
+    # Succeed
+    zarrurl = str(tmp_path / "D.zarr")
+    da.ones(shape=(8, 8)).to_zarr(f"{zarrurl}/0")
+    build_pyramid(zarrurl=zarrurl, coarsening_xy=2, num_levels=3)
+    # Should fail because overwrite is not set
+    with pytest.raises(ValueError):
+        build_pyramid(
+            zarrurl=zarrurl, coarsening_xy=2, num_levels=3, overwrite=False
+        )
+    # Should work
+    build_pyramid(
+        zarrurl=zarrurl, coarsening_xy=2, num_levels=3, overwrite=True
+    )