make DataSaver picklable

python-adaptive · Apr 9, 2020 · cb2cb25 · cb2cb25
1 parent aff10f3
commit cb2cb25
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/adaptive/learner/data_saver.py b/adaptive/learner/data_saver.py
@@ -51,6 +51,18 @@ def _set_data(self, data):
         learner_data, self.extra_data = data
         self.learner._set_data(learner_data)
 
+    def __getstate__(self):
+        return (
+            self.learner,
+            self.arg_picker,
+            self._get_data(),
+        )
+
+    def __setstate__(self, state):
+        learner, arg_picker, data = state
+        self.__init__(learner, arg_picker)
+        self._set_data(data)
+
     @copy_docstring_from(BaseLearner.save)
     def save(self, fname, compress=True):
         # We copy this method because the 'DataSaver' is not a