type hint fixes for adaptive/learner/balancing_learner.py

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -9,20 +9,12 @@
 import numpy as np
 from numpy import float64, int64
 
-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.learner.sequence_learner import SequenceLearner, _IgnoreFirstArgument
 from adaptive.notebook_integration import ensure_holoviews
 from adaptive.utils import cache_latest, named_product, restore
 
 
-def dispatch(
-    child_functions: Union[List[Callable], List[partial], List[_IgnoreFirstArgument]],
-    arg: Any,
-) -> Union[int, float64, float]:
+def dispatch(child_functions: List[Callable], arg: Any,) -> Union[int, float64, float]:
     index, x = arg
     return child_functions[index](x)
 
@@ -79,17 +71,7 @@ class BalancingLearner(BaseLearner):
     """
 
     def __init__(
-        self,
-        learners: Union[
-            List[SequenceLearner],
-            List[AverageLearner],
-            List[Learner2D],
-            List[Learner1D],
-            List[LearnerND],
-        ],
-        *,
-        cdims=None,
-        strategy="loss_improvements"
+        self, learners: List[BaseLearner], *, cdims=None, strategy="loss_improvements"
     ) -> None:
         self.learners = learners
 
@@ -246,7 +228,7 @@ def _ask_and_tell_based_on_cycle(
 
         return points, loss_improvements
 
-    def ask(self, n: int, tell_pending: bool = True) -> Any:
+    def ask(self, n: int, tell_pending: bool = True) -> Tuple[List[Any], List[float]]:
         """Chose points for learners."""
         if n == 0:
             return [], []
@@ -369,7 +351,9 @@ def remove_unfinished(self):
             learner.remove_unfinished()
 
     @classmethod
-    def from_product(cls, f, learner_type, learner_kwargs, combos):
+    def from_product(
+        cls, f, learner_type, learner_kwargs, combos
+    ) -> "BalancingLearner":
         """Create a `BalancingLearner` with learners of all combinations of
         named variables’ values. The `cdims` will be set correctly, so calling
         `learner.plot` will be a `holoviews.core.HoloMap` with the correct labels.