add type annotations for adaptive/runner.py

Author: basnijholt

adaptive/runner.py

@@ -8,8 +8,20 @@
 import time
 import traceback
 import warnings
+from _asyncio import Future, Task
+from concurrent.futures.process import ProcessPoolExecutor
 from contextlib import suppress
+from typing import Any, Callable, List, Optional, Set, Tuple, Union
 
+from distributed.cfexecutor import ClientExecutor
+from distributed.client import Client
+from ipyparallel.client.asyncresult import AsyncResult
+from ipyparallel.client.view import ViewExecutor
+from numpy import float64
+
+from adaptive.learner.learner1D import Learner1D
+from adaptive.learner.learner2D import Learner2D
+from adaptive.learner.learnerND import LearnerND
 from adaptive.notebook_integration import in_ipynb, live_info, live_plot
 
 try:
@@ -121,16 +133,16 @@ class BaseRunner(metaclass=abc.ABCMeta):
 
     def __init__(
         self,
-        learner,
-        goal,
+        learner: Union[Learner1D, Learner2D, LearnerND],
+        goal: Callable,
         *,
         executor=None,
         ntasks=None,
         log=False,
         shutdown_executor=False,
         retries=0,
         raise_if_retries_exceeded=True,
-    ):
+    ) -> None:
 
         self.executor = _ensure_executor(executor)
         self.goal = goal
@@ -157,7 +169,7 @@ def __init__(
         self.to_retry = {}
         self.tracebacks = {}
 
-    def _get_max_tasks(self):
+    def _get_max_tasks(self) -> int:
         return self._max_tasks or _get_ncores(self.executor)
 
     def _do_raise(self, e, x):
@@ -169,10 +181,10 @@ def _do_raise(self, e, x):
         ) from e
 
     @property
-    def do_log(self):
+    def do_log(self) -> bool:
         return self.log is not None
 
-    def _ask(self, n):
+    def _ask(self, n: int) -> Any:
         points = [
             p for p in self.to_retry.keys() if p not in self.pending_points.values()
         ][:n]
@@ -206,7 +218,9 @@ def overhead(self):
         t_total = self.elapsed_time()
         return (1 - t_function / t_total) * 100
 
-    def _process_futures(self, done_futs):
+    def _process_futures(
+        self, done_futs: Union[Set[Future], Set[Future], Set[AsyncResult], Set[Task]]
+    ) -> None:
         for fut in done_futs:
             x = self.pending_points.pop(fut)
             try:
@@ -227,7 +241,9 @@ def _process_futures(self, done_futs):
                     self.log.append(("tell", x, y))
                 self.learner.tell(x, y)
 
-    def _get_futures(self):
+    def _get_futures(
+        self,
+    ) -> Union[List[Task], List[Future], List[Future], List[AsyncResult]]:
         # Launch tasks to replace the ones that completed
         # on the last iteration, making sure to fill workers
         # that have started since the last iteration.
@@ -248,7 +264,7 @@ def _get_futures(self):
         futures = list(self.pending_points.keys())
         return futures
 
-    def _remove_unfinished(self):
+    def _remove_unfinished(self) -> List[Future]:
         # remove points with 'None' values from the learner
         self.learner.remove_unfinished()
         # cancel any outstanding tasks
@@ -257,7 +273,7 @@ def _remove_unfinished(self):
             fut.cancel()
         return remaining
 
-    def _cleanup(self):
+    def _cleanup(self) -> None:
         if self.shutdown_executor:
             # XXX: temporary set wait=True for Python 3.7
             # see https://github.com/python-adaptive/adaptive/issues/156
@@ -347,16 +363,16 @@ class BlockingRunner(BaseRunner):
 
     def __init__(
         self,
-        learner,
-        goal,
+        learner: Union[LearnerND, Learner2D, Learner1D],
+        goal: Callable,
         *,
         executor=None,
         ntasks=None,
         log=False,
         shutdown_executor=False,
         retries=0,
         raise_if_retries_exceeded=True,
-    ):
+    ) -> None:
         if inspect.iscoroutinefunction(learner.function):
             raise ValueError(
                 "Coroutine functions can only be used " "with 'AsyncRunner'."
@@ -373,10 +389,12 @@ def __init__(
         )
         self._run()
 
-    def _submit(self, x):
+    def _submit(
+        self, x: Union[Tuple[int, int], int, Tuple[float64, float64], float]
+    ) -> Union[Future, AsyncResult]:
         return self.executor.submit(self.learner.function, x)
 
-    def _run(self):
+    def _run(self) -> None:
         first_completed = concurrent.FIRST_COMPLETED
 
         if self._get_max_tasks() < 1:
@@ -476,8 +494,8 @@ class AsyncRunner(BaseRunner):
 
     def __init__(
         self,
-        learner,
-        goal=None,
+        learner: Union[Learner1D, Learner2D],
+        goal: Optional[Callable] = None,
         *,
         executor=None,
         ntasks=None,
@@ -486,7 +504,7 @@ def __init__(
         ioloop=None,
         retries=0,
         raise_if_retries_exceeded=True,
-    ):
+    ) -> None:
 
         if goal is None:
 
@@ -539,7 +557,9 @@ def goal(_):
                 "'adaptive.notebook_extension()'"
             )
 
-    def _submit(self, x):
+    def _submit(
+        self, x: Union[Tuple[int, int], int, Tuple[float64, float64], float]
+    ) -> Union[Task, Future]:
         ioloop = self.ioloop
         if inspect.iscoroutinefunction(self.learner.function):
             return ioloop.create_task(self.learner.function(x))
@@ -604,7 +624,7 @@ def live_info(self, *, update_interval=0.1):
         """
         return live_info(self, update_interval=update_interval)
 
-    async def _run(self):
+    async def _run(self) -> None:
         first_completed = asyncio.FIRST_COMPLETED
 
         if self._get_max_tasks() < 1:
@@ -668,7 +688,7 @@ async def _saver(save_kwargs=save_kwargs, interval=interval):
 Runner = AsyncRunner
 
 
-def simple(learner, goal):
+def simple(learner: Any, goal: Callable) -> None:
     """Run the learner until the goal is reached.
 
     Requests a single point from the learner, evaluates
@@ -694,7 +714,16 @@ def simple(learner, goal):
             learner.tell(x, y)
 
 
-def replay_log(learner, log):
+def replay_log(
+    learner: LearnerND,
+    log: List[
+        Union[
+            Tuple[str, int],
+            Tuple[str, Tuple[int, int, int], float],
+            Tuple[str, Tuple[float, float, float], float],
+        ]
+    ],
+) -> None:
     """Apply a sequence of method calls to a learner.
 
     This is useful for debugging runners.
@@ -713,7 +742,7 @@ def replay_log(learner, log):
 # --- Useful runner goals
 
 
-def stop_after(*, seconds=0, minutes=0, hours=0):
+def stop_after(*, seconds=0, minutes=0, hours=0) -> Callable:
     """Stop a runner after a specified time.
 
     For example, to specify a runner that should stop after
@@ -756,7 +785,7 @@ class SequentialExecutor(concurrent.Executor):
     This executor is mainly for testing.
     """
 
-    def submit(self, fn, *args, **kwargs):
+    def submit(self, fn: Callable, *args, **kwargs) -> Future:
         fut = concurrent.Future()
         try:
             fut.set_result(fn(*args, **kwargs))
@@ -771,7 +800,9 @@ def shutdown(self, wait=True):
         pass
 
 
-def _ensure_executor(executor):
+def _ensure_executor(
+    executor: Optional[Union[Client, Client, ProcessPoolExecutor, SequentialExecutor]]
+) -> Union[SequentialExecutor, ProcessPoolExecutor, ViewExecutor, ClientExecutor]:
     if executor is None:
         executor = _default_executor(**_default_executor_kwargs)
 
@@ -788,7 +819,9 @@ def _ensure_executor(executor):
         )
 
 
-def _get_ncores(ex):
+def _get_ncores(
+    ex: Union[SequentialExecutor, ProcessPoolExecutor, ViewExecutor, ClientExecutor]
+) -> int:
     """Return the maximum  number of cores that an executor can use."""
     if with_ipyparallel and isinstance(ex, ipyparallel.client.view.ViewExecutor):
         return len(ex.view)