Handle errors when saving files gracefully

Fixes jazzband#118.

* Prevent existing files from being wiped when writing fails
* Backup previous file if it exists when writing succeeds

Signed-off-by: Christopher Arndt <[email protected]>

CHANGELOG.md

@@ -13,7 +13,11 @@ This document records all notable changes to Watson. This project adheres to
   timespan to the current year, month, week or day.
 * Added: Zsh completion support
 * Added: document installation via homebrew on OS X
+* Updated: when saving the Watson frames, state or config file, the most recent
+  previous version of the file is kept as a back up.
 * Fixed: bash completion of projects and tags with spaces in them
+* Fixed: if saving the Watson frames, state or config file fails for any
+  reason, the original is kept (and not wiped as before).
 
 ## 1.3.2 (2016-03-01)
 

tests/test_watson.py

@@ -9,9 +9,9 @@
     import mock
 
 try:
-    from io import StringIO
-except ImportError:
     from StringIO import StringIO
+except ImportError:
+    from io import StringIO
 
 import py
 import pytest
@@ -23,7 +23,8 @@
 from click import get_app_dir
 from watson import Watson, WatsonError
 from watson.watson import ConfigurationError, ConfigParser
-from watson.utils import get_start_time_for_period
+from watson.utils import get_start_time_for_period, make_json_writer, safe_save
+
 
 TEST_FIXTURE_DIR = py.path.local(
     os.path.dirname(
@@ -630,6 +631,20 @@ def test_save_empty_last_sync(config_dir):
             assert json_mock.call_args[0][0] == 0
 
 
+def test_watson_save_calls_safe_save(watson, config_dir):
+    frames_file = os.path.join(config_dir, 'frames')
+    watson.start('foo', tags=['A', 'B'])
+    watson.stop()
+
+    with mock.patch('watson.watson.safe_save') as save_mock:
+        watson.save()
+
+        assert watson._frames.changed
+        assert save_mock.call_count == 1
+        assert len(save_mock.call_args[0]) == 2
+        assert save_mock.call_args[0][0] == frames_file
+
+
 # push
 
 def test_push_with_no_config(watson):
@@ -881,3 +896,78 @@ def test_merge_report(watson, datafiles):
 def test_get_start_time_for_period(now, mode, start_time):
     with mock_datetime(now, datetime):
         assert get_start_time_for_period(mode).datetime == start_time
+
+
+# utils
+
+def test_make_json_writer():
+    fp = StringIO()
+    writer = make_json_writer(lambda: {'foo': 42})
+    writer(fp)
+    assert fp.getvalue() == '{\n "foo": 42\n}'
+
+
+def test_make_json_writer_with_args():
+    fp = StringIO()
+    writer = make_json_writer(lambda x: {'foo': x}, 23)
+    writer(fp)
+    assert fp.getvalue() == '{\n "foo": 23\n}'
+
+
+def test_make_json_writer_with_kwargs():
+    fp = StringIO()
+    writer = make_json_writer(lambda foo=None: {'foo': foo}, foo='bar')
+    writer(fp)
+    assert fp.getvalue() == '{\n "foo": "bar"\n}'
+
+
+def test_safe_save(config_dir):
+    save_file = os.path.join(config_dir, 'test')
+    backup_file = os.path.join(config_dir, 'test' + '.bak')
+
+    assert not os.path.exists(save_file)
+    safe_save(save_file, lambda f: f.write("Success"))
+    assert os.path.exists(save_file)
+    assert not os.path.exists(backup_file)
+
+    with open(save_file) as fp:
+        assert fp.read() == "Success"
+
+    safe_save(save_file, "Again")
+    assert os.path.exists(backup_file)
+
+    with open(save_file) as fp:
+        assert fp.read() == "Again"
+
+    with open(backup_file) as fp:
+        assert fp.read() == "Success"
+
+    assert os.path.getmtime(save_file) >= os.path.getmtime(backup_file)
+
+
+def test_safe_save_with_exception(config_dir):
+    save_file = os.path.join(config_dir, 'test')
+    backup_file = os.path.join(config_dir, 'test' + '.bak')
+
+    def failing_writer(f):
+        raise RuntimeError("Save failed.")
+
+    assert not os.path.exists(save_file)
+
+    with pytest.raises(RuntimeError):
+        safe_save(save_file, failing_writer)
+
+    assert not os.path.exists(save_file)
+    assert not os.path.exists(backup_file)
+
+    safe_save(save_file, lambda f: f.write("Success"))
+    assert os.path.exists(save_file)
+    assert not os.path.exists(backup_file)
+
+    with pytest.raises(RuntimeError):
+        safe_save(save_file, failing_writer)
+
+    with open(save_file) as fp:
+        assert fp.read() == "Success"
+
+    assert not os.path.exists(backup_file)

watson/cli.py

@@ -12,10 +12,11 @@
 import arrow
 import click
 
-from . import watson
+from . import watson as _watson
 from .frames import Frame
-from .utils import (format_timedelta, get_frame_from_argument, options,
-                    sorted_groupby, style, get_start_time_for_period)
+from .utils import (format_timedelta, get_frame_from_argument,
+                    get_start_time_for_period, options, safe_save,
+                    sorted_groupby, style)
 
 
 class MutuallyExclusiveOption(click.Option):
@@ -43,7 +44,7 @@ def format_message(self):
         return style('error', self.message)
 
 
-watson.WatsonError = WatsonCliError
+_watson.WatsonError = WatsonCliError
 
 
 class DateParamType(click.ParamType):
@@ -63,7 +64,7 @@ def convert(self, value, param, ctx):
 
 
 @click.group()
-@click.version_option(version=watson.__version__, prog_name='Watson')
+@click.version_option(version=_watson.__version__, prog_name='Watson')
 @click.pass_context
 def cli(ctx):
     """
@@ -76,7 +77,7 @@ def cli(ctx):
 
     # This is the main command group, needed by click in order
     # to handle the subcommands
-    ctx.obj = watson.Watson(config_dir=os.environ.get('WATSON_DIR'))
+    ctx.obj = _watson.Watson(config_dir=os.environ.get('WATSON_DIR'))
 
 
 @cli.command()
@@ -811,15 +812,19 @@ def config(context, key, value, edit):
     config = watson.config
 
     if edit:
-        click.edit(filename=watson.config_file, extension='.ini')
+        with open(watson.config_file) as fp:
+            newconfig = click.edit(text=fp.read(), extension='.ini')
+
+        if newconfig:
+            safe_save(watson.config_file, newconfig)
 
         try:
             watson.config = None
             watson.config
-        except WatsonCliError:
+        except _watson.ConfigurationError as exc:
             watson.config = config
             watson.save()
-            raise
+            raise WatsonCliError(str(exc))
         return
 
     if not key:

watson/utils.py

@@ -1,12 +1,21 @@
-import itertools
 import datetime
+import itertools
+import json
+import os
+import tempfile
 
 import click
 import arrow
 
 from click.exceptions import UsageError
 
 
+try:
+    text_type = (str, unicode)
+except NameError:
+    text_type = str
+
+
 def style(name, element):
     def _style_tags(tags):
         if not tags:
@@ -138,3 +147,53 @@ def get_start_time_for_period(period):
         raise ValueError('Unsupported period value: {}'.format(period))
 
     return start_time
+
+
+def make_json_writer(func, *args, **kwargs):
+    """
+    Return a function that receives a file-like object and writes the return
+    value of func(*args, **kwargs) as JSON to it.
+    """
+    def writer(f):
+        json.dump(func(*args, **kwargs), f, indent=1, ensure_ascii=False)
+    return writer
+
+
+def safe_save(path, content, ext='.bak'):
+    """
+    Save given content to file at given path safely.
+
+    `content` may either be a (unicode) string to write to the file, or a
+    function taking one argument, a file object opened for writing. The
+    function may write (unicode) strings to the file object (but doesn't need
+    to close it).
+
+    The file to write to is created at a temporary location first. If there is
+    an error creating or writing to the temp file or calling `content`, the
+    destination file is left untouched. Otherwise, if all is well, an existing
+    destination file is backed up to `path` + `ext` (defaults to '.bak') and
+    the temporary file moved into its place.
+
+    """
+    tmpfp = tempfile.NamedTemporaryFile(mode='w+', delete=False)
+    try:
+        with tmpfp:
+            if isinstance(content, text_type):
+                tmpfp.write(content)
+            else:
+                content(tmpfp)
+    except:
+        try:
+            os.unlink(tmpfp.name)
+        except (IOError, OSError):
+            pass
+        raise
+    else:
+        if os.path.exists(path):
+            try:
+                os.unlink(path + ext)
+            except OSError:
+                pass
+            os.rename(path, path + ext)
+
+        os.rename(tmpfp.name, path)

watson/watson.py

@@ -14,14 +14,15 @@
 
 from .config import ConfigParser
 from .frames import Frames
+from .utils import make_json_writer, safe_save
 from .version import version as __version__  # noqa
 
 
 class WatsonError(RuntimeError):
     pass
 
 
-class ConfigurationError(WatsonError, configparser.Error):
+class ConfigurationError(configparser.Error, WatsonError):
     pass
 
 
@@ -147,21 +148,18 @@ def save(self):
                 else:
                     current = {}
 
-                with open(self.state_file, 'w+') as f:
-                    json.dump(current, f, indent=1, ensure_ascii=False)
+                safe_save(self.state_file, make_json_writer(lambda: current))
 
             if self._frames is not None and self._frames.changed:
-                with open(self.frames_file, 'w+') as f:
-                    json.dump(self.frames.dump(), f, indent=1,
-                              ensure_ascii=False)
+                safe_save(self.frames_file,
+                          make_json_writer(self.frames.dump))
 
             if self._config_changed:
-                with open(self.config_file, 'w+') as f:
-                    self.config.write(f)
+                safe_save(self.config_file, self.config.write)
 
             if self._last_sync is not None:
-                with open(self.last_sync_file, 'w+') as f:
-                    json.dump(self._format_date(self.last_sync), f)
+                safe_save(self.last_sync_file,
+                          make_json_writer(self._format_date, self.last_sync))
         except OSError as e:
             raise WatsonError(
                 "Impossible to write {}: {}".format(e.filename, e)