add type annotations for adaptive/tests/test_runner.py

Author: basnijholt

adaptive/tests/test_runner.py

@@ -1,7 +1,9 @@
 import asyncio
 import time
+from typing import Callable, Iterator, Union
 
 import pytest
+from distributed.client import Client
 
 from adaptive.learner import Learner1D, Learner2D
 from adaptive.runner import (
@@ -15,24 +17,24 @@
 )
 
 
-def blocking_runner(learner, goal):
+def blocking_runner(learner: Union[Learner1D, Learner2D], goal: Callable) -> None:
     BlockingRunner(learner, goal, executor=SequentialExecutor())
 
 
-def async_runner(learner, goal):
+def async_runner(learner: Union[Learner1D, Learner2D], goal: Callable) -> None:
     runner = AsyncRunner(learner, goal, executor=SequentialExecutor())
     asyncio.get_event_loop().run_until_complete(runner.task)
 
 
 runners = [simple, blocking_runner, async_runner]
 
 
-def trivial_goal(learner):
+def trivial_goal(learner: Union[Learner1D, Learner2D]) -> bool:
     return learner.npoints > 10
 
 
 @pytest.mark.parametrize("runner", runners)
-def test_simple(runner):
+def test_simple(runner: Callable) -> None:
     """Test that the runners actually run."""
 
     def f(x):
@@ -44,7 +46,7 @@ def f(x):
 
 
 @pytest.mark.parametrize("runner", runners)
-def test_nonconforming_output(runner):
+def test_nonconforming_output(runner: Callable) -> None:
     """Test that using a runner works with a 2D learner, even when the
     learned function outputs a 1-vector. This tests against the regression
     flagged in https://github.com/python-adaptive/adaptive/issues/81.
@@ -56,7 +58,7 @@ def f(x):
     runner(Learner2D(f, [(-1, 1), (-1, 1)]), trivial_goal)
 
 
-def test_aync_def_function():
+def test_aync_def_function() -> None:
     async def f(x):
         return x
 
@@ -69,7 +71,7 @@ async def f(x):
 
 
 @pytest.fixture(scope="session")
-def ipyparallel_executor():
+def ipyparallel_executor() -> Iterator[Client]:
     from ipyparallel import Client
     import pexpect
 
@@ -81,7 +83,7 @@ def ipyparallel_executor():
 
 
 @pytest.fixture(scope="session")
-def dask_executor():
+def dask_executor() -> Iterator[Client]:
     from distributed import Client
 
     client = Client(n_workers=1)
@@ -93,7 +95,7 @@ def linear(x):
     return x
 
 
-def test_concurrent_futures_executor():
+def test_concurrent_futures_executor() -> None:
     from concurrent.futures import ProcessPoolExecutor
 
     BlockingRunner(
@@ -103,7 +105,7 @@ def test_concurrent_futures_executor():
     )
 
 
-def test_stop_after_goal():
+def test_stop_after_goal() -> None:
     seconds_to_wait = 0.2  # don't make this too large or the test will take ages
     start_time = time.time()
     BlockingRunner(Learner1D(linear, (-1, 1)), stop_after(seconds=seconds_to_wait))
@@ -112,15 +114,15 @@ def test_stop_after_goal():
 
 
 @pytest.mark.skipif(not with_ipyparallel, reason="IPyparallel is not installed")
-def test_ipyparallel_executor(ipyparallel_executor):
+def test_ipyparallel_executor(ipyparallel_executor: Client) -> None:
     learner = Learner1D(linear, (-1, 1))
     BlockingRunner(learner, trivial_goal, executor=ipyparallel_executor)
     assert learner.npoints > 0
 
 
 @pytest.mark.timeout(60)
 @pytest.mark.skipif(not with_distributed, reason="dask.distributed is not installed")
-def test_distributed_executor(dask_executor):
+def test_distributed_executor(dask_executor: Client) -> None:
     learner = Learner1D(linear, (-1, 1))
     BlockingRunner(learner, trivial_goal, executor=dask_executor)
     assert learner.npoints > 0