use BaseLearner instead of Union of some learners

Author: basnijholt

adaptive/learner/data_saver.py

@@ -3,11 +3,7 @@
 from operator import itemgetter
 from typing import Callable, Dict, Tuple, Union
 
-from adaptive.learner.average_learner import AverageLearner
 from adaptive.learner.base_learner import BaseLearner
-from adaptive.learner.learner1D import Learner1D
-from adaptive.learner.learner2D import Learner2D
-from adaptive.learner.learnerND import LearnerND
 from adaptive.utils import copy_docstring_from
 
 
@@ -31,11 +27,7 @@ class DataSaver:
     >>> learner = DataSaver(_learner, arg_picker=itemgetter('y'))
     """
 
-    def __init__(
-        self,
-        learner: Union[Learner2D, Learner1D, LearnerND, AverageLearner],
-        arg_picker: itemgetter,
-    ) -> None:
+    def __init__(self, learner: BaseLearner, arg_picker: itemgetter,) -> None:
         self.learner = learner
         self.extra_data = OrderedDict()
         self.function = learner.function

adaptive/runner.py

@@ -19,9 +19,7 @@
 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.learner import BaseLearner
 from adaptive.notebook_integration import in_ipynb, live_info, live_plot
 
 try:
@@ -133,7 +131,7 @@ class BaseRunner(metaclass=abc.ABCMeta):
 
     def __init__(
         self,
-        learner: Union[Learner1D, Learner2D, LearnerND],
+        learner: BaseLearner,
         goal: Callable,
         *,
         executor=None,
@@ -363,7 +361,7 @@ class BlockingRunner(BaseRunner):
 
     def __init__(
         self,
-        learner: Union[LearnerND, Learner2D, Learner1D],
+        learner: BaseLearner,
         goal: Callable,
         *,
         executor=None,
@@ -494,7 +492,7 @@ class AsyncRunner(BaseRunner):
 
     def __init__(
         self,
-        learner: Union[Learner1D, Learner2D],
+        learner: BaseLearner,
         goal: Optional[Callable] = None,
         *,
         executor=None,
@@ -688,7 +686,7 @@ async def _saver(save_kwargs=save_kwargs, interval=interval):
 Runner = AsyncRunner
 
 
-def simple(learner: Any, goal: Callable) -> None:
+def simple(learner: BaseLearner, goal: Callable) -> None:
     """Run the learner until the goal is reached.
 
     Requests a single point from the learner, evaluates
@@ -715,7 +713,7 @@ def simple(learner: Any, goal: Callable) -> None:
 
 
 def replay_log(
-    learner: LearnerND,
+    learner: BaseLearner,
     log: List[
         Union[
             Tuple[str, int],

adaptive/tests/test_runner.py

@@ -1,11 +1,11 @@
 import asyncio
 import time
-from typing import Callable, Iterator, Union
+from typing import Callable, Iterator
 
 import pytest
 from distributed.client import Client
 
-from adaptive.learner import Learner1D, Learner2D
+from adaptive.learner import BaseLearner, Learner1D, Learner2D
 from adaptive.runner import (
     AsyncRunner,
     BlockingRunner,
@@ -17,19 +17,19 @@
 )
 
 
-def blocking_runner(learner: Union[Learner1D, Learner2D], goal: Callable) -> None:
+def blocking_runner(learner: BaseLearner, goal: Callable) -> None:
     BlockingRunner(learner, goal, executor=SequentialExecutor())
 
 
-def async_runner(learner: Union[Learner1D, Learner2D], goal: Callable) -> None:
+def async_runner(learner: BaseLearner, 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: Union[Learner1D, Learner2D]) -> bool:
+def trivial_goal(learner: BaseLearner) -> bool:
     return learner.npoints > 10