allow RemoteDataset as argument (#1382)

* allow RemoteDataset as argument

* add changelog

---------

Co-authored-by: valentin-pinkau <[email protected]>

Author: valentin-pinkau

webknossos/Changelog.md

@@ -37,6 +37,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 - The returned dataset ID from dataset exploration is now used. [#1378](https://github.com/scalableminds/webknossos-libs/pull/1378) 
 - Refactored the architecture, by introducing RemoteLayers, RemoteSegmentationLayers and their abstract base classes. [#1371](https://github.com/scalableminds/webknossos-libs/pull/1371])
 - Updated the api version of the webknossos-api to 12. [#1371](https://github.com/scalableminds/webknossos-libs/pull/1371])
+- Allowing RemoteDataset to align mags, when down- or upsampling [#1382](https://github.com/scalableminds/webknossos-libs/pull/1382)
 
 ### Fixed
 

webknossos/webknossos/dataset/layer/_downsampling_utils.py

@@ -4,13 +4,13 @@
 from collections.abc import Callable
 from enum import Enum
 from itertools import product
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Union
 
 import numpy as np
 from scipy.ndimage import zoom
 
 if TYPE_CHECKING:
-    from webknossos.dataset.dataset import Dataset
+    from webknossos.dataset.dataset import Dataset, RemoteDataset
 
 from webknossos.dataset_properties import LayerCategoryType
 from webknossos.geometry import Mag, Vec3Int, Vec3IntLike
@@ -46,7 +46,7 @@ def determine_upsample_buffer_shape(array_info: ArrayInfo) -> Vec3Int:
 def calculate_mags_to_downsample(
     from_mag: Mag,
     coarsest_mag: Mag,
-    dataset_to_align_with: Optional["Dataset"],
+    dataset_to_align_with: Union["Dataset", "RemoteDataset"] | None,
     voxel_size: tuple[float, float, float] | None,
 ) -> list[Mag]:
     assert np.all(from_mag.to_np() <= coarsest_mag.to_np())
@@ -119,7 +119,7 @@ def calculate_mags_to_downsample(
 def calculate_mags_to_upsample(
     from_mag: Mag,
     finest_mag: Mag,
-    dataset_to_align_with: Optional["Dataset"],
+    dataset_to_align_with: Union["Dataset", "RemoteDataset"] | None,
     voxel_size: tuple[float, float, float] | None,
 ) -> list[Mag]:
     return list(

webknossos/webknossos/dataset/layer/layer.py

@@ -3,7 +3,7 @@
 from inspect import getframeinfo, stack
 from os import PathLike
 from os.path import relpath
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Union
 
 import numpy as np
 from cluster_tools import Executor
@@ -39,7 +39,7 @@
 )
 
 if TYPE_CHECKING:
-    from webknossos.dataset import Dataset
+    from webknossos.dataset import Dataset, RemoteDataset
     from webknossos.dataset.layer import RemoteLayer
 
     from .segmentation_layer import SegmentationLayer
@@ -955,8 +955,8 @@ def _create_dir_for_mag(self, mag: MagLike) -> UPath:
         return full_path
 
     def _get_dataset_from_align_with_other_layers(
-        self, align_with_other_layers: Union[bool, "Dataset"]
-    ) -> Optional["Dataset"]:
+        self, align_with_other_layers: Union[bool, "Dataset", "RemoteDataset"]
+    ) -> Union["Dataset", "RemoteDataset"] | None:
         if isinstance(align_with_other_layers, bool):
             return self.dataset if align_with_other_layers else None
         else:
@@ -970,7 +970,7 @@ def downsample(
         interpolation_mode: str = "default",
         compress: bool | Zarr3Config = True,
         sampling_mode: str | SamplingModes = SamplingModes.ANISOTROPIC,
-        align_with_other_layers: Union[bool, "Dataset"] = True,
+        align_with_other_layers: Union[bool, "Dataset", "RemoteDataset"] = True,
         buffer_shape: Vec3Int | None = None,
         force_sampling_scheme: bool = False,
         allow_overwrite: bool = False,
@@ -990,7 +990,7 @@ def downsample(
             compress (bool | Zarr3Config): Whether to compress the generated magnifications. For Zarr3 datasets, codec configuration and chunk key encoding may also be supplied. Defaults to True.
             sampling_mode (str | SamplingModes): How dimensions should be downsampled.
                 Defaults to ANISOTROPIC.
-            align_with_other_layers (bool | Dataset): Whether to align with other layers. True by default.
+            align_with_other_layers (bool | Dataset | RemoteDataset): Whether to align with other layers. True by default.
             buffer_shape (Vec3Int | None): Shape of processing buffer. Defaults to None.
             force_sampling_scheme (bool): Force invalid sampling schemes. Defaults to False.
             allow_overwrite (bool): Whether existing mags can be overwritten. False by default.