run test_average_learner.py without type warnings

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -371,7 +371,7 @@ def from_product(
         f,
         learner_type: BaseLearner,
         learner_kwargs: Dict[str, Any],
-        combos: Dict[str, Iterable[Any]],
+        combos: Dict[str, Sequence[Any]],
     ) -> "BalancingLearner":
         """Create a `BalancingLearner` with learners of all combinations of
         named variables’ values. The `cdims` will be set correctly, so calling

adaptive/learner/learnerND.py

@@ -383,7 +383,7 @@ def npoints(self) -> int:
         return len(self.data)
 
     @property
-    def vdim(self):
+    def vdim(self) -> int:
         """Length of the output of ``learner.function``.
         If the output is unsized (when it's a scalar)
         then `vdim = 1`.
@@ -399,10 +399,10 @@ def vdim(self):
         return self._vdim if self._vdim is not None else 1
 
     @property
-    def bounds_are_done(self):
+    def bounds_are_done(self) -> bool:
         return all(p in self.data for p in self._bounds_points)
 
-    def _ip(self):
+    def _ip(self) -> interpolate.LinearNDInterpolator:
         """A `scipy.interpolate.LinearNDInterpolator` instance
         containing the learner's data."""
         # XXX: take our own triangulation into account when generating the _ip
@@ -427,7 +427,7 @@ def tri(self) -> Optional[Triangulation]:
         return self._tri
 
     @property
-    def values(self):
+    def values(self) -> np.ndarray:
         """Get the values from `data` as a numpy array."""
         return np.array(list(self.data.values()), dtype=float)
 

adaptive/runner.py

@@ -237,20 +237,20 @@ def __init__(
         # Error handling attributes
         self.retries = retries
         self.raise_if_retries_exceeded = raise_if_retries_exceeded
-        self._to_retry = {}
-        self._tracebacks = {}
+        self._to_retry: Dict[int, int] = {}
+        self._tracebacks: Dict[int, str] = {}
 
-        self._id_to_point = {}
+        self._id_to_point: Dict[int, Any] = {}
         self._next_id = functools.partial(
             next, itertools.count()
         )  # some unique id to be associated with each point
 
     def _get_max_tasks(self) -> int:
         return self._max_tasks or _get_ncores(self.executor)
 
-    def _do_raise(self, e, i):
-        tb = self._tracebacks[i]
-        x = self._id_to_point[i]
+    def _do_raise(self, e, pid):
+        tb = self._tracebacks[pid]
+        x = self._id_to_point[pid]
         raise RuntimeError(
             "An error occured while evaluating "
             f'"learner.function({x})". '
@@ -379,7 +379,7 @@ def _cleanup(self) -> None:
         self.end_time = time.time()
 
     @property
-    def failed(self):
+    def failed(self) -> Set[Any]:
         """Set of points that failed ``runner.retries`` times."""
         return set(self._tracebacks) - set(self._to_retry)
 
@@ -398,15 +398,15 @@ def _submit(self, x):
         pass
 
     @property
-    def tracebacks(self):
+    def tracebacks(self) -> List[Tuple[int, str]]:
         return [(self._id_to_point[pid], tb) for pid, tb in self._tracebacks.items()]
 
     @property
-    def to_retry(self):
+    def to_retry(self) -> List[Tuple[int, int]]:
         return [(self._id_to_point[pid], n) for pid, n in self._to_retry.items()]
 
     @property
-    def pending_points(self):
+    def pending_points(self) -> List[Tuple[Future, Any]]:
         return [
             (fut, self._id_to_point[pid]) for fut, pid in self._pending_tasks.items()
         ]

adaptive/utils.py

@@ -5,12 +5,12 @@
 import pickle
 from contextlib import contextmanager
 from itertools import product
-from typing import Any, Callable, Dict, Iterable, Iterator
+from typing import Any, Callable, Dict, Iterator, Sequence
 
 from atomicwrites import AtomicWriter
 
 
-def named_product(**items: Dict[str, Iterable[Any]]):
+def named_product(**items: Dict[str, Sequence[Any]]):
     names = items.keys()
     vals = items.values()
     return [dict(zip(names, res)) for res in product(*vals)]