refactor Executor type to protocol

cluster_tools/Changelog.md

@@ -16,6 +16,7 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 - Added `DaskScheduler` (only Python >= 3.9). [#943](https://github.com/scalableminds/webknossos-libs/pull/943)
 
 ### Changed
+- The exported `Executor` type is now implemented as a protocol.  [#943](https://github.com/scalableminds/webknossos-libs/pull/943)
 
 ### Fixed
 

cluster_tools/cluster_tools/__init__.py

@@ -1,5 +1,6 @@
-from typing import Any, Literal, Union, overload
+from typing import Any, Literal, overload
 
+from cluster_tools.executor_protocol import Executor
 from cluster_tools.executors.dask import DaskExecutor
 from cluster_tools.executors.debug_sequential import DebugSequentialExecutor
 from cluster_tools.executors.multiprocessing_ import MultiprocessingExecutor
@@ -128,6 +129,3 @@ def get_executor(environment: str, **kwargs: Any) -> "Executor":
     elif environment == "test_pickling":
         return PickleExecutor(**kwargs)
     raise Exception("Unknown executor: {}".format(environment))
-
-
-Executor = Union[ClusterExecutor, MultiprocessingExecutor, DaskExecutor]

cluster_tools/cluster_tools/executor_protocol.py

@@ -0,0 +1,50 @@
+from concurrent.futures import Future
+from os import PathLike
+from typing import Callable, Iterable, Iterator, List, Optional, Protocol, TypeVar
+
+from typing_extensions import ParamSpec
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+_S = TypeVar("_S")
+
+
+class Executor(Protocol):
+    @classmethod
+    def as_completed(cls, futures: List["Future[_T]"]) -> Iterator["Future[_T]"]:
+        ...
+
+    def submit(
+        self,
+        __fn: Callable[_P, _T],
+        /,
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> "Future[_T]":
+        ...
+
+    def map_unordered(self, fn: Callable[[_S], _T], args: Iterable[_S]) -> Iterator[_T]:
+        ...
+
+    def map_to_futures(
+        self,
+        fn: Callable[[_S], _T],
+        args: Iterable[_S],
+        output_pickle_path_getter: Optional[Callable[[_S], PathLike]],
+    ) -> List["Future[_T]"]:
+        ...
+
+    def map(
+        self,
+        fn: Callable[[_S], _T],
+        iterables: Iterable[_S],
+        timeout: Optional[float],
+        chunksize: int,
+    ) -> Iterator[_T]:
+        ...
+
+    def forward_log(self, fut: "Future[_T]") -> _T:
+        ...
+
+    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
+        ...

cluster_tools/cluster_tools/executors/dask.py

@@ -107,15 +107,15 @@ def map_to_futures(
 
         return futs
 
-    def map(
+    def map(  # type: ignore[override]
         self,
         fn: Callable[[_S], _T],
-        *iterables: Iterable[Any],
+        iterables: Iterable[Any],
         timeout: Optional[float] = None,
         chunksize: int = 1,
     ) -> Iterator[_T]:
         return iter(
-            list(super().map(fn, *iterables, timeout=timeout, chunksize=chunksize))
+            list(super().map(fn, [iterables], timeout=timeout, chunksize=chunksize))
         )
 
     def forward_log(self, fut: "Future[_T]") -> _T:

cluster_tools/cluster_tools/executors/multiprocessing_.py

@@ -148,6 +148,15 @@ def submit(  # type: ignore[override]
         enrich_future_with_uncaught_warning(fut)
         return fut
 
+    def map(  # type: ignore[override]
+        self,
+        fn: Callable[[_S], _T],
+        iterables: Iterable[Any],
+        timeout: Optional[float] = None,
+        chunksize: int = 1,
+    ) -> Iterator[_T]:
+        return super().map(fn, [iterables], timeout=timeout, chunksize=chunksize)
+
     def _submit_via_io(
         self,
         __fn: Callable[_P, _T],

cluster_tools/cluster_tools/executors/pickle_.py

@@ -1,19 +1,28 @@
 from concurrent.futures import Future
-from typing import Any, Callable, TypeVar
+from functools import partial
+from typing import Callable, TypeVar
+
+from typing_extensions import ParamSpec
 
 from cluster_tools._utils import pickling
 from cluster_tools.executors.multiprocessing_ import MultiprocessingExecutor
 
 # The module name includes a _-suffix to avoid name clashes with the standard library pickle module.
 
 _T = TypeVar("_T")
+_P = ParamSpec("_P")
+_S = TypeVar("_S")
 
 
-def _pickle_identity(obj: _T) -> _T:
+def _pickle_identity(obj: _S) -> _S:
     return pickling.loads(pickling.dumps(obj))
 
 
-def _pickle_identity_executor(fn: Callable[..., _T], *args: Any, **kwargs: Any) -> _T:
+def _pickle_identity_executor(
+    fn: Callable[_P, _T],
+    *args: _P.args,
+    **kwargs: _P.kwargs,
+) -> _T:
     result = fn(*args, **kwargs)
     return _pickle_identity(result)
 
@@ -27,13 +36,16 @@ class PickleExecutor(MultiprocessingExecutor):
 
     def submit(  # type: ignore[override]
         self,
-        fn: Callable[..., _T],
-        *args: Any,
-        **kwargs: Any,
+        fn: Callable[_P, _T],
+        /,
+        *args: _P.args,
+        **kwargs: _P.kwargs,
     ) -> "Future[_T]":
         (fn_pickled, args_pickled, kwargs_pickled) = _pickle_identity(
             (fn, args, kwargs)
         )
         return super().submit(
-            _pickle_identity_executor, fn_pickled, *args_pickled, **kwargs_pickled
+            partial(_pickle_identity_executor, fn_pickled),
+            *args_pickled,
+            **kwargs_pickled,
         )

cluster_tools/cluster_tools/schedulers/cluster_executor.py

@@ -598,10 +598,9 @@ def shutdown(self, wait: bool = True, cancel_futures: bool = True) -> None:
                 pass
         self.files_to_clean_up = []
 
-    # TODO: args should be *iterables, this would be a breaking change!
     def map(  # type: ignore[override]
         self,
-        fn: Callable[_P, _T],
+        fn: Callable[[_S], _T],
         args: Iterable[Any],
         timeout: Optional[float] = None,
         chunksize: Optional[int] = None,

webknossos/Changelog.md

@@ -13,6 +13,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 [Commits](https://github.com/scalableminds/webknossos-libs/compare/v0.13.6...HEAD)
 
 ### Breaking Changes
+- `wait_and_ensure_success` from `webknossos.utils` now requires a `executor` argument. [#943](https://github.com/scalableminds/webknossos-libs/pull/943)
 
 ### Added