ensure atomic writes when saving a file

Right now, if a program crashes in the middle of saving, you lose
all your data.

Author: basnijholt

adaptive/learner/base_learner.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 
 import abc
-import os
 from contextlib import suppress
 from copy import deepcopy
 
@@ -175,16 +174,8 @@ def save(self, fname, compress=True):
             using compression, one must load it with compression too.
         """
         data = self._get_data()
-
-        backup = os.path.join(".", fname)
-        if os.path.exists(fname):
-            os.rename(backup)
-
         save(fname, data, compress)
 
-        if os.path.exists(backup):
-            os.remove(backup)
-
     def load(self, fname, compress=True):
         """Load the data of a learner from a pickle file.
 
@@ -196,12 +187,6 @@ def load(self, fname, compress=True):
             If the data is compressed when saved, one must load it
             with compression too.
         """
-        with suppress(FileNotFoundError):
-            backup = os.path.join(".", fname)
-            if os.path.getctime(fname) < os.path.getctime(backup):
-                # the backup is the newer file
-                fname = backup
-
         with suppress(FileNotFoundError, EOFError):
             data = load(fname, compress)
             self._set_data(data)

adaptive/utils.py

@@ -7,6 +7,7 @@
 from contextlib import contextmanager
 from itertools import product
 
+from atomicwrites import AtomicWriter
 
 def named_product(**items):
     names = items.keys()
@@ -44,10 +45,13 @@ def save(fname, data, compress=True):
     dirname = os.path.dirname(fname)
     if dirname:
         os.makedirs(dirname, exist_ok=True)
-    _open = gzip.open if compress else open
-    with _open(fname, "wb") as f:
-        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
 
+    blob = pickle.dumps(data, protocol=pickle.HIGHEST_PROTOCOL)
+    if compress:
+        blob = gzip.compress(blob)
+
+    with AtomicWriter(fname, 'wb', overwrite=True).open() as f:
+        f.write(blob)
 
 def load(fname, compress=True):
     fname = os.path.expanduser(fname)

docs/environment.yml

@@ -12,6 +12,7 @@ dependencies:
   - bokeh=1.0.4
   - plotly=3.9.0
   - ipywidgets=7.4.2
+  - atomicwrites=1.3.0
   - pip:
     - git+https://github.com/basnijholt/jupyter-sphinx.git@widgets_execute
     - sphinx_fontawesome==0.0.6

environment.yml

@@ -16,3 +16,4 @@ dependencies:
   - ipywidgets
   - scikit-optimize
   - plotly
+  - atomicwrites

setup.py

@@ -24,7 +24,7 @@ def get_version_and_cmdclass(package_name):
 version, cmdclass = get_version_and_cmdclass("adaptive")
 
 
-install_requires = ["scipy", "sortedcollections >= 1.1", "sortedcontainers >= 2.0"]
+install_requires = ["scipy", "sortedcollections >= 1.1", "sortedcontainers >= 2.0", "atomicwrites"]
 
 extras_require = {
     "notebook": [