Add mag from existing zarr array (#1151)

* Add first implementation for

* Update changelog.

Author: markbader

webknossos/Changelog.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Breaking Changes
 
 ### Added
+- Added `add_mag_from_zarrarray` to `Layer` class, to add existing Zarr arrays as a mag of a layer. [#1151](https://github.com/scalableminds/webknossos-libs/pull/1151)
 
 ### Changed
 

webknossos/examples/add_existing_zarr_array.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+
+import webknossos as wk
+
+OUTPUT_PATH = Path(__file__).parent.parent / "testoutput"
+ARRAY_PATH = (
+    Path(__file__).parent.parent / "testdata" / "simple_zarr3_dataset" / "color" / "1"
+)
+
+
+def main() -> None:
+    ds = wk.Dataset(OUTPUT_PATH, voxel_size=(10, 10, 10))
+    layer = ds.add_layer("color", category="color", data_format="zarr3")
+
+    layer.add_mag_from_zarrarray(1, ARRAY_PATH)
+
+
+if __name__ == "__main__":
+    main()

webknossos/webknossos/dataset/layer.py

@@ -14,7 +14,7 @@
 from upath import UPath
 
 from ..geometry import Mag, NDBoundingBox, Vec3Int, Vec3IntLike
-from ._array import ArrayException, BaseArray, DataFormat
+from ._array import ArrayException, BaseArray, DataFormat, ZarritaArray
 from ._downsampling_utils import (
     calculate_default_coarsest_mag,
     calculate_mags_to_downsample,
@@ -44,6 +44,7 @@
     get_executor_for_args,
     is_fs_path,
     is_remote_path,
+    movetree,
     named_partial,
     rmtree,
     strip_trailing_slash,
@@ -753,6 +754,50 @@ def add_fs_copy_mag(
 
         return mag
 
+    def add_mag_from_zarrarray(
+        self,
+        mag: Union[int, str, list, tuple, np.ndarray, Mag],
+        path: PathLike,
+        move: bool = False,
+        extend_layer_bounding_box: bool = True,
+    ) -> MagView:
+        """
+        Copies the data at `path` to the current layer of the dataset
+        via the filesystem and adds it as `mag`. When `move` flag is set
+        the array is moved, otherwise a copy of the zarrarray is created.
+        """
+        self.dataset._ensure_writable()
+        source_path = Path(path)
+
+        try:
+            ZarritaArray.open(source_path)
+        except ArrayException as e:
+            raise ValueError(
+                "The given path does not lead to a valid Zarr Array: "
+            ) from e
+        else:
+            mag = Mag(mag)
+            self._assert_mag_does_not_exist_yet(mag)
+            if move:
+                movetree(source_path, self.path / str(mag))
+            else:
+                copytree(source_path, self.path / str(mag))
+
+            mag_view = self.add_mag_for_existing_files(mag)
+
+            if extend_layer_bounding_box:
+                # assumption: the topleft of the bounding box is still the same, the size might differ
+                # axes of the layer and the zarr array provided are the same
+                zarray_size = (
+                    mag_view.info.shape[mag_view.info.dimension_names.index(axis)]
+                    for axis in self.bounding_box.axes
+                    if axis != "c"
+                )
+                size = self.bounding_box.size.pairmax(zarray_size)
+                self.bounding_box = self.bounding_box.with_size(size)
+
+            return mag_view
+
     def _create_dir_for_mag(
         self, mag: Union[int, str, list, tuple, np.ndarray, Mag]
     ) -> None:

webknossos/webknossos/utils.py

@@ -13,7 +13,7 @@
 from multiprocessing import cpu_count
 from os.path import relpath
 from pathlib import Path
-from shutil import copyfileobj
+from shutil import copyfileobj, move
 from typing import (
     Any,
     Callable,
@@ -346,6 +346,10 @@ def _append(path: Path, parts: Tuple[str, ...]) -> Path:
                 copyfileobj(in_file, out_file)
 
 
+def movetree(in_path: Path, out_path: Path) -> None:
+    move(in_path, out_path)
+
+
 K = TypeVar("K")  # key
 V = TypeVar("V")  # value
 C = TypeVar("C")  # cache