add type annotations for adaptive/learner/data_saver.py

Author: basnijholt

adaptive/learner/data_saver.py

@@ -1,7 +1,13 @@
 import functools
 from collections import OrderedDict
+from operator import itemgetter
+from typing import Callable, Dict, Tuple, Union
 
+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.utils import copy_docstring_from
 
 
@@ -25,13 +31,17 @@ class DataSaver:
     >>> learner = DataSaver(_learner, arg_picker=itemgetter('y'))
     """
 
-    def __init__(self, learner, arg_picker):
+    def __init__(
+        self,
+        learner: Union[Learner2D, Learner1D, LearnerND, AverageLearner],
+        arg_picker: itemgetter,
+    ) -> None:
         self.learner = learner
         self.extra_data = OrderedDict()
         self.function = learner.function
         self.arg_picker = arg_picker
 
-    def __getattr__(self, attr):
+    def __getattr__(self, attr: str) -> Union[Callable, int]:
         return getattr(self.learner, attr)
 
     @copy_docstring_from(BaseLearner.tell)
@@ -44,10 +54,23 @@ def tell(self, x, result):
     def tell_pending(self, x):
         self.learner.tell_pending(x)
 
-    def _get_data(self):
+    def _get_data(
+        self,
+    ) -> Union[
+        Tuple[Dict[Union[int, float], float], OrderedDict],
+        Tuple[OrderedDict, OrderedDict],
+        Tuple[Tuple[Dict[int, float], int, float, float], OrderedDict],
+    ]:
         return self.learner._get_data(), self.extra_data
 
-    def _set_data(self, data):
+    def _set_data(
+        self,
+        data: Union[
+            Tuple[OrderedDict, OrderedDict],
+            Tuple[Dict[Union[int, float], float], OrderedDict],
+            Tuple[Tuple[Dict[int, float], int, float, float], OrderedDict],
+        ],
+    ) -> None:
         learner_data, self.extra_data = data
         self.learner._set_data(learner_data)