more type fixes

Author: basnijholt

adaptive/learner/base_learner.py

@@ -118,7 +118,7 @@ def remove_unfinished(self):
         pass
 
     @abc.abstractmethod
-    def loss(self, real=True):
+    def loss(self, real: bool = True):
         """Return the loss for the current state of the learner.
 
         Parameters
@@ -130,7 +130,7 @@ def loss(self, real=True):
         """
 
     @abc.abstractmethod
-    def ask(self, n, tell_pending=True):
+    def ask(self, n: int, tell_pending: bool = True):
         """Choose the next 'n' points to evaluate.
 
         Parameters

adaptive/learner/learner2D.py

@@ -598,7 +598,9 @@ def _fill_stack(
 
         return points_new, losses_new
 
-    def ask(self, n: int, tell_pending: bool = True) -> Tuple[np.ndarray, np.ndarray]:
+    def ask(
+        self, n: int, tell_pending: bool = True
+    ) -> Tuple[List[Tuple[float, float]], List[float]]:
         # Even if tell_pending is False we add the point such that _fill_stack
         # will return new points, later we remove these points if needed.
         points = list(self._stack.keys())

adaptive/learner/learnerND.py

@@ -31,7 +31,7 @@ def to_list(inp: float) -> List[float]:
     return [inp]
 
 
-def volume(simplex: List[Tuple[float, float]], ys: None = None) -> float:
+def volume(simplex: Simplex, ys: None = None) -> float:
     # Notice the parameter ys is there so you can use this volume method as
     # as loss function
     matrix = np.subtract(simplex[:-1], simplex[-1], dtype=float)
@@ -69,7 +69,7 @@ def uniform_loss(simplex: np.ndarray, values: np.ndarray, value_scale: float) ->
     return volume(simplex)
 
 
-def std_loss(simplex: np.ndarray, values: np.ndarray, value_scale: float) -> np.ndarray:
+def std_loss(simplex: Simplex, values: np.ndarray, value_scale: float) -> np.ndarray:
     """
     Computes the loss of the simplex based on the standard deviation.
 

adaptive/learner/triangulation.py

@@ -52,7 +52,7 @@ def fast_norm(v: Union[Tuple[float, ...], ndarray]) -> float:
 
 def fast_2d_point_in_simplex(
     point: Point, simplex: SimplexPoints, eps: float = 1e-8
-) -> Union[np.bool_, bool]:
+) -> Union[bool, np.bool_]:
     (p0x, p0y), (p1x, p1y), (p2x, p2y) = simplex
     px, py = point
 
@@ -68,7 +68,7 @@ def fast_2d_point_in_simplex(
 
 def point_in_simplex(
     point: Point, simplex: SimplexPoints, eps: float = 1e-8
-) -> Union[np.bool_, bool]:
+) -> Union[bool, np.bool_]:
     if len(point) == 2:
         return fast_2d_point_in_simplex(point, simplex, eps)
 
@@ -417,7 +417,7 @@ def get_reduced_simplex(
 
     def point_in_simplex(
         self, point: Point, simplex: Simplex, eps: float = 1e-8
-    ) -> Union[np.bool_, bool]:
+    ) -> Union[bool, np.bool_]:
         vertices = self.get_vertices(simplex)
         return point_in_simplex(point, vertices, eps)
 

adaptive/utils.py

@@ -3,9 +3,9 @@
 import gzip
 import os
 import pickle
-from contextlib import contextmanager
+from contextlib import _GeneratorContextManager, contextmanager
 from itertools import product
-from typing import Any, Callable, Dict, Iterator, Sequence
+from typing import Any, Callable, Dict, Sequence
 
 from atomicwrites import AtomicWriter
 
@@ -17,7 +17,7 @@ def named_product(**items: Dict[str, Sequence[Any]]):
 
 
 @contextmanager
-def restore(*learners) -> Iterator[None]:
+def restore(*learners) -> _GeneratorContextManager:
     states = [learner.__getstate__() for learner in learners]
     try:
         yield