add type annotations for adaptive/learner/sequence_learner.py

Author: basnijholt

adaptive/learner/sequence_learner.py

@@ -1,5 +1,8 @@
 from copy import copy
+from functools import partial
+from typing import Any, List, Tuple, Union
 
+from numpy import float64, ndarray
 from sortedcontainers import SortedDict, SortedSet
 
 from adaptive.learner.base_learner import BaseLearner
@@ -15,17 +18,22 @@ class _IgnoreFirstArgument:
     pickable.
     """
 
-    def __init__(self, function):
+    def __init__(self, function: partial) -> None:
         self.function = function
 
-    def __call__(self, index_point, *args, **kwargs):
+    def __call__(
+        self,
+        index_point: Union[Tuple[int, int], Tuple[int, float64], Tuple[int, ndarray]],
+        *args,
+        **kwargs
+    ) -> Union[float64, float]:
         index, point = index_point
         return self.function(point, *args, **kwargs)
 
-    def __getstate__(self):
+    def __getstate__(self) -> partial:
         return self.function
 
-    def __setstate__(self, function):
+    def __setstate__(self, function: partial) -> None:
         self.__init__(function)
 
 
@@ -56,7 +64,7 @@ class SequenceLearner(BaseLearner):
     the added benefit of having results in the local kernel already.
     """
 
-    def __init__(self, function, sequence):
+    def __init__(self, function: partial, sequence: Union[range, ndarray]) -> None:
         self._original_function = function
         self.function = _IgnoreFirstArgument(function)
         self._to_do_indices = SortedSet({i for i, _ in enumerate(sequence)})
@@ -65,7 +73,13 @@ def __init__(self, function, sequence):
         self.data = SortedDict()
         self.pending_points = set()
 
-    def ask(self, n, tell_pending=True):
+    def ask(
+        self, n: int, tell_pending: bool = True
+    ) -> Union[
+        Tuple[List[Tuple[int, float64]], List[float]],
+        Tuple[List[Tuple[int, int]], List[float]],
+        Tuple[List[Tuple[int, ndarray]], List[float]],
+    ]:
         indices = []
         points = []
         loss_improvements = []
@@ -83,17 +97,17 @@ def ask(self, n, tell_pending=True):
 
         return points, loss_improvements
 
-    def _get_data(self):
+    def _get_data(self) -> SortedDict:
         return self.data
 
-    def _set_data(self, data):
+    def _set_data(self, data: SortedDict) -> None:
         if data:
             indices, values = zip(*data.items())
             # the points aren't used by tell, so we can safely pass None
             points = [(i, None) for i in indices]
             self.tell_many(points, values)
 
-    def loss(self, real=True):
+    def loss(self, real: bool = True) -> float:
         if not (self._to_do_indices or self.pending_points):
             return 0
         else:
@@ -105,13 +119,19 @@ def remove_unfinished(self):
             self._to_do_indices.add(i)
         self.pending_points = set()
 
-    def tell(self, point, value):
+    def tell(
+        self,
+        point: Union[
+            Tuple[int, int], Tuple[int, float64], Tuple[int, ndarray], Tuple[int, None]
+        ],
+        value: Union[float64, float],
+    ) -> None:
         index, point = point
         self.data[index] = value
         self.pending_points.discard(index)
         self._to_do_indices.discard(index)
 
-    def tell_pending(self, point):
+    def tell_pending(self, point: Any) -> None:
         index, point = point
         self.pending_points.add(index)
         self._to_do_indices.discard(index)
@@ -126,5 +146,5 @@ def result(self):
         return list(self.data.values())
 
     @property
-    def npoints(self):
+    def npoints(self) -> int:
         return len(self.data)