move to Layer.add_mag_as_copy

Author: normanrz

webknossos/tests/dataset/test_dataset.py

@@ -5,6 +5,7 @@
 from collections.abc import Iterator
 from pathlib import Path
 from typing import cast
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -2113,7 +2114,15 @@ def test_add_fs_copy_mag(data_format: DataFormat, output_path: UPath) -> None:
     copy_layer = copy_ds.add_layer(
         "color", COLOR_CATEGORY, dtype_per_channel="uint8", data_format=data_format
     )
-    copy_mag = copy_layer.add_fs_copy_mag(original_mag, extend_layer_bounding_box=True)
+
+    with mock.patch.object(
+        copy_layer, "add_fs_copy_mag", wraps=copy_layer.add_fs_copy_mag
+    ) as mocked_method:
+        copy_mag = copy_layer.add_mag_as_copy(
+            original_mag, extend_layer_bounding_box=True
+        )
+        mocked_method.assert_called_once()
+
     assert not copy_layer.read_only
     assert not copy_mag.read_only
 

webknossos/webknossos/dataset/dataset.py

@@ -40,7 +40,7 @@
 )
 from ..geometry.mag import MagLike
 from ..geometry.nd_bounding_box import derive_nd_bounding_box_from_shape
-from ._array import ArrayException, ArrayInfo, BaseArray, Zarr3ArrayInfo, Zarr3Config
+from ._array import ArrayException, ArrayInfo, BaseArray, Zarr3Config
 from ._metadata import DatasetMetadata
 from ._utils import pims_images
 from .defaults import (
@@ -2308,48 +2308,17 @@ def add_layer_as_copy(
                 f"Copying {mag_view.layer.name}/{mag_view.mag.to_layer_name()}"
             )
 
-            can_use_fs_copy = True
-            if (
-                (chunk_shape is not None and chunk_shape != mag_view.info.chunk_shape)
-                or (
-                    shard_shape is not None and shard_shape != mag_view.info.shard_shape
-                )
-                or (
-                    chunks_per_shard is not None
-                    and chunks_per_shard != mag_view.info.chunks_per_shard
-                )
-                or (
-                    data_format is not None and data_format != mag_view.info.data_format
-                )
-                or (
-                    compress is not None
-                    and (
-                        isinstance(compress, Zarr3Config)
-                        and isinstance(mag_view.info, Zarr3ArrayInfo)
-                        and compress != mag_view.info.zarr3_config
-                    )
-                    or compress != mag_view.info.compression_mode
-                )
-            ):
-                can_use_fs_copy = False
-
-            if can_use_fs_copy:
-                layer.add_fs_copy_mag(
-                    mag_view,
-                    extend_layer_bounding_box=False,
-                )
-            else:
-                layer.add_mag_as_copy(
-                    mag_view,
-                    extend_layer_bounding_box=False,
-                    chunk_shape=chunk_shape,
-                    shard_shape=shard_shape,
-                    chunks_per_shard=chunks_per_shard,
-                    compress=compress,
-                    exists_ok=exists_ok,
-                    executor=executor,
-                    progress_desc=progress_desc,
-                )
+            layer.add_mag_as_copy(
+                mag_view,
+                extend_layer_bounding_box=False,
+                chunk_shape=chunk_shape,
+                shard_shape=shard_shape,
+                chunks_per_shard=chunks_per_shard,
+                compress=compress,
+                exists_ok=exists_ok,
+                executor=executor,
+                progress_desc=progress_desc,
+            )
 
         if (
             with_attachments

webknossos/webknossos/dataset/layer.py

@@ -14,7 +14,7 @@
 
 from ..geometry import Mag, NDBoundingBox, Vec3Int, Vec3IntLike
 from ..geometry.mag import MagLike
-from ._array import ArrayException, TensorStoreArray, Zarr3Config
+from ._array import ArrayException, TensorStoreArray, Zarr3ArrayInfo, Zarr3Config
 from ._downsampling_utils import (
     calculate_default_coarsest_mag,
     calculate_mags_to_downsample,
@@ -60,6 +60,8 @@
     DEFAULT_SHARD_SHAPE,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def _is_int(s: str) -> bool:
     try:
@@ -870,6 +872,44 @@ def add_mag_as_copy(
         self._ensure_writable()
         foreign_mag_view = MagView._ensure_mag_view(foreign_mag_view_or_path)
 
+        if (
+            (
+                chunk_shape is None
+                or Vec3Int.from_vec_or_int(chunk_shape)
+                == foreign_mag_view.info.chunk_shape
+            )
+            and (
+                shard_shape is None
+                or Vec3Int.from_vec_or_int(shard_shape)
+                == foreign_mag_view.info.shard_shape
+            )
+            and (
+                chunks_per_shard is None
+                or Vec3Int.from_vec_or_int(chunks_per_shard)
+                == foreign_mag_view.info.chunks_per_shard
+            )
+            and (self.data_format == foreign_mag_view.info.data_format)
+            and (
+                compress is None
+                or (
+                    (
+                        not isinstance(compress, Zarr3Config)
+                        or not isinstance(foreign_mag_view.info, Zarr3ArrayInfo)
+                        or compress == foreign_mag_view.info.zarr3_config
+                    )
+                    and compress == foreign_mag_view.info.compression_mode
+                )
+            )
+        ):
+            logger.debug(
+                f"Optimization: Copying files from {foreign_mag_view.path} to {self.path}/{foreign_mag_view.mag} directly without re-encoding."
+            )
+            return self.add_fs_copy_mag(
+                foreign_mag_view,
+                extend_layer_bounding_box=extend_layer_bounding_box,
+                exists_ok=exists_ok,
+            )
+
         chunk_shape = Vec3Int.from_vec_or_int(
             chunk_shape or foreign_mag_view.info.chunk_shape
         )
@@ -1041,6 +1081,7 @@ def add_fs_copy_mag(
         foreign_mag_view_or_path: PathLike | str | MagView,
         *,
         extend_layer_bounding_box: bool = True,
+        exists_ok: bool = False,
     ) -> MagView:
         """
         Copies the data at `foreign_mag_view_or_path` which belongs to another dataset to the current dataset via the filesystem.
@@ -1050,7 +1091,12 @@ def add_fs_copy_mag(
         foreign_mag_view = MagView._ensure_mag_view(foreign_mag_view_or_path)
         self._assert_mag_does_not_exist_yet(foreign_mag_view.mag)
 
+        assert foreign_mag_view.info.data_format == self.data_format
+
         mag_path = self.path / str(foreign_mag_view.mag)
+        assert exists_ok or not mag_path.exists(), (
+            f"Cannot copy {foreign_mag_view.path} to {mag_path} because it already exists and `exist_ok` is set to False."
+        )
         copytree(
             foreign_mag_view.path,
             mag_path,