type hint fixes for adaptive/learner/sequence_learner.py

Author: basnijholt

adaptive/learner/sequence_learner.py

@@ -1,6 +1,5 @@
 from copy import copy
-from functools import partial
-from typing import Any, List, Tuple, Union
+from typing import Any, Callable, List, Sequence, Tuple, Union
 
 import numpy as np
 from sortedcontainers import SortedDict, SortedSet
@@ -18,22 +17,19 @@ class _IgnoreFirstArgument:
     pickable.
     """
 
-    def __init__(self, function: partial) -> None:
+    def __init__(self, function: Callable) -> None:
         self.function = function
 
     def __call__(
-        self,
-        index_point: Union[Tuple[int, int], Tuple[int, float], Tuple[int, np.ndarray]],
-        *args,
-        **kwargs
+        self, index_point: Tuple[int, Union[float, np.ndarray]], *args, **kwargs
     ) -> float:
         index, point = index_point
         return self.function(point, *args, **kwargs)
 
-    def __getstate__(self) -> partial:
+    def __getstate__(self) -> Callable:
         return self.function
 
-    def __setstate__(self, function: partial) -> None:
+    def __setstate__(self, function: Callable) -> None:
         self.__init__(function)
 
 
@@ -64,7 +60,7 @@ class SequenceLearner(BaseLearner):
     the added benefit of having results in the local kernel already.
     """
 
-    def __init__(self, function: partial, sequence: Union[range, np.ndarray]) -> None:
+    def __init__(self, function: Callable, sequence: Sequence) -> None:
         self._original_function = function
         self.function = _IgnoreFirstArgument(function)
         self._to_do_indices = SortedSet({i for i, _ in enumerate(sequence)})
@@ -73,13 +69,7 @@ def __init__(self, function: partial, sequence: Union[range, np.ndarray]) -> Non
         self.data = SortedDict()
         self.pending_points = set()
 
-    def ask(
-        self, n: int, tell_pending: bool = True
-    ) -> Union[
-        Tuple[List[Tuple[int, float]], List[float]],
-        Tuple[List[Tuple[int, int]], List[float]],
-        Tuple[List[Tuple[int, np.ndarray]], List[float]],
-    ]:
+    def ask(self, n: int, tell_pending: bool = True) -> Tuple[Any, List[float]]:
         indices = []
         points = []
         loss_improvements = []
@@ -119,16 +109,7 @@ def remove_unfinished(self):
             self._to_do_indices.add(i)
         self.pending_points = set()
 
-    def tell(
-        self,
-        point: Union[
-            Tuple[int, int],
-            Tuple[int, float],
-            Tuple[int, np.ndarray],
-            Tuple[int, None],
-        ],
-        value: float,
-    ) -> None:
+    def tell(self, point: Tuple[int, Any], value: Any,) -> None:
         index, point = point
         self.data[index] = value
         self.pending_points.discard(index)