Skip to content

Commit 6712b3a

Browse files
committed
move to Layer.add_mag_as_copy
1 parent d8a0572 commit 6712b3a

File tree

3 files changed

+69
-45
lines changed

3 files changed

+69
-45
lines changed

webknossos/tests/dataset/test_dataset.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from collections.abc import Iterator
66
from pathlib import Path
77
from typing import cast
8+
from unittest import mock
89

910
import numpy as np
1011
import pytest
@@ -2113,7 +2114,15 @@ def test_add_fs_copy_mag(data_format: DataFormat, output_path: UPath) -> None:
21132114
copy_layer = copy_ds.add_layer(
21142115
"color", COLOR_CATEGORY, dtype_per_channel="uint8", data_format=data_format
21152116
)
2116-
copy_mag = copy_layer.add_fs_copy_mag(original_mag, extend_layer_bounding_box=True)
2117+
2118+
with mock.patch.object(
2119+
copy_layer, "add_fs_copy_mag", wraps=copy_layer.add_fs_copy_mag
2120+
) as mocked_method:
2121+
copy_mag = copy_layer.add_mag_as_copy(
2122+
original_mag, extend_layer_bounding_box=True
2123+
)
2124+
mocked_method.assert_called_once()
2125+
21172126
assert not copy_layer.read_only
21182127
assert not copy_mag.read_only
21192128

webknossos/webknossos/dataset/dataset.py

Lines changed: 12 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
)
4141
from ..geometry.mag import MagLike
4242
from ..geometry.nd_bounding_box import derive_nd_bounding_box_from_shape
43-
from ._array import ArrayException, ArrayInfo, BaseArray, Zarr3ArrayInfo, Zarr3Config
43+
from ._array import ArrayException, ArrayInfo, BaseArray, Zarr3Config
4444
from ._metadata import DatasetMetadata
4545
from ._utils import pims_images
4646
from .defaults import (
@@ -2308,48 +2308,17 @@ def add_layer_as_copy(
23082308
f"Copying {mag_view.layer.name}/{mag_view.mag.to_layer_name()}"
23092309
)
23102310

2311-
can_use_fs_copy = True
2312-
if (
2313-
(chunk_shape is not None and chunk_shape != mag_view.info.chunk_shape)
2314-
or (
2315-
shard_shape is not None and shard_shape != mag_view.info.shard_shape
2316-
)
2317-
or (
2318-
chunks_per_shard is not None
2319-
and chunks_per_shard != mag_view.info.chunks_per_shard
2320-
)
2321-
or (
2322-
data_format is not None and data_format != mag_view.info.data_format
2323-
)
2324-
or (
2325-
compress is not None
2326-
and (
2327-
isinstance(compress, Zarr3Config)
2328-
and isinstance(mag_view.info, Zarr3ArrayInfo)
2329-
and compress != mag_view.info.zarr3_config
2330-
)
2331-
or compress != mag_view.info.compression_mode
2332-
)
2333-
):
2334-
can_use_fs_copy = False
2335-
2336-
if can_use_fs_copy:
2337-
layer.add_fs_copy_mag(
2338-
mag_view,
2339-
extend_layer_bounding_box=False,
2340-
)
2341-
else:
2342-
layer.add_mag_as_copy(
2343-
mag_view,
2344-
extend_layer_bounding_box=False,
2345-
chunk_shape=chunk_shape,
2346-
shard_shape=shard_shape,
2347-
chunks_per_shard=chunks_per_shard,
2348-
compress=compress,
2349-
exists_ok=exists_ok,
2350-
executor=executor,
2351-
progress_desc=progress_desc,
2352-
)
2311+
layer.add_mag_as_copy(
2312+
mag_view,
2313+
extend_layer_bounding_box=False,
2314+
chunk_shape=chunk_shape,
2315+
shard_shape=shard_shape,
2316+
chunks_per_shard=chunks_per_shard,
2317+
compress=compress,
2318+
exists_ok=exists_ok,
2319+
executor=executor,
2320+
progress_desc=progress_desc,
2321+
)
23532322

23542323
if (
23552324
with_attachments

webknossos/webknossos/dataset/layer.py

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
from ..geometry import Mag, NDBoundingBox, Vec3Int, Vec3IntLike
1616
from ..geometry.mag import MagLike
17-
from ._array import ArrayException, TensorStoreArray, Zarr3Config
17+
from ._array import ArrayException, TensorStoreArray, Zarr3ArrayInfo, Zarr3Config
1818
from ._downsampling_utils import (
1919
calculate_default_coarsest_mag,
2020
calculate_mags_to_downsample,
@@ -60,6 +60,8 @@
6060
DEFAULT_SHARD_SHAPE,
6161
)
6262

63+
logger = logging.getLogger(__name__)
64+
6365

6466
def _is_int(s: str) -> bool:
6567
try:
@@ -870,6 +872,44 @@ def add_mag_as_copy(
870872
self._ensure_writable()
871873
foreign_mag_view = MagView._ensure_mag_view(foreign_mag_view_or_path)
872874

875+
if (
876+
(
877+
chunk_shape is None
878+
or Vec3Int.from_vec_or_int(chunk_shape)
879+
== foreign_mag_view.info.chunk_shape
880+
)
881+
and (
882+
shard_shape is None
883+
or Vec3Int.from_vec_or_int(shard_shape)
884+
== foreign_mag_view.info.shard_shape
885+
)
886+
and (
887+
chunks_per_shard is None
888+
or Vec3Int.from_vec_or_int(chunks_per_shard)
889+
== foreign_mag_view.info.chunks_per_shard
890+
)
891+
and (self.data_format == foreign_mag_view.info.data_format)
892+
and (
893+
compress is None
894+
or (
895+
(
896+
not isinstance(compress, Zarr3Config)
897+
or not isinstance(foreign_mag_view.info, Zarr3ArrayInfo)
898+
or compress == foreign_mag_view.info.zarr3_config
899+
)
900+
and compress == foreign_mag_view.info.compression_mode
901+
)
902+
)
903+
):
904+
logger.debug(
905+
f"Optimization: Copying files from {foreign_mag_view.path} to {self.path}/{foreign_mag_view.mag} directly without re-encoding."
906+
)
907+
return self.add_fs_copy_mag(
908+
foreign_mag_view,
909+
extend_layer_bounding_box=extend_layer_bounding_box,
910+
exists_ok=exists_ok,
911+
)
912+
873913
chunk_shape = Vec3Int.from_vec_or_int(
874914
chunk_shape or foreign_mag_view.info.chunk_shape
875915
)
@@ -1041,6 +1081,7 @@ def add_fs_copy_mag(
10411081
foreign_mag_view_or_path: PathLike | str | MagView,
10421082
*,
10431083
extend_layer_bounding_box: bool = True,
1084+
exists_ok: bool = False,
10441085
) -> MagView:
10451086
"""
10461087
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(
10501091
foreign_mag_view = MagView._ensure_mag_view(foreign_mag_view_or_path)
10511092
self._assert_mag_does_not_exist_yet(foreign_mag_view.mag)
10521093

1094+
assert foreign_mag_view.info.data_format == self.data_format
1095+
10531096
mag_path = self.path / str(foreign_mag_view.mag)
1097+
assert exists_ok or not mag_path.exists(), (
1098+
f"Cannot copy {foreign_mag_view.path} to {mag_path} because it already exists and `exist_ok` is set to False."
1099+
)
10541100
copytree(
10551101
foreign_mag_view.path,
10561102
mag_path,

0 commit comments

Comments
 (0)