Merge pull request #57 from thewtex/cache-cleanup

ENH: Clean up scale cache stores before process exit

ngff_zarr/ngff_image.py

@@ -1,10 +1,12 @@
-from dataclasses import dataclass
-from typing import Dict, Mapping, Optional, Sequence
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Mapping, Optional, Sequence
 
 from dask.array.core import Array as DaskArray
 
 from .zarr_metadata import Units
 
+ComputedCallback = Callable[[], None]
+
 
 @dataclass
 class NgffImage:
@@ -14,3 +16,4 @@ class NgffImage:
     translation: Dict[str, float]
     name: str = "image"
     axes_units: Optional[Mapping[str, Units]] = None
+    computed_callbacks: List[ComputedCallback] = field(default_factory=list)

ngff_zarr/to_multiscales.py

@@ -91,6 +91,7 @@ def remove_from_cache_store(sig_id, frame):  # noqa: ARG001
     atexit.register(remove_from_cache_store, None, None)
     signal.signal(signal.SIGTERM, remove_from_cache_store)
     signal.signal(signal.SIGINT, remove_from_cache_store)
+    image.computed_callbacks.append(lambda: remove_from_cache_store(None, None))
 
     data = image.data
 

ngff_zarr/to_ngff_zarr.py

@@ -258,6 +258,9 @@ def to_ngff_zarr(
             and multiscales.method
             and multiscales.chunks
         ):
+            for callback in image.computed_callbacks:
+                callback()
+
             image.data = dask.array.from_zarr(store, component=path)
             next_multiscales_factor = multiscales.scale_factors[index]
             if isinstance(next_multiscales_factor, int):
@@ -285,4 +288,8 @@ def to_ngff_zarr(
         elif index < nscales - 1:
             next_image = multiscales.images[index + 1]
 
+    for image in multiscales.images:
+        for callback in image.computed_callbacks:
+            callback()
+
     zarr.consolidate_metadata(store)