diff --git a/IPython/html/services/contents/filemanager.py b/IPython/html/services/contents/filemanager.py index ad36f39..ef6cbdf 100644 --- a/IPython/html/services/contents/filemanager.py +++ b/IPython/html/services/contents/filemanager.py @@ -313,7 +313,7 @@ class FileContentsManager(ContentsManager): bcontent = base64.decodestring(b64_bytes) except Exception as e: raise web.HTTPError(400, u'Encoding error saving %s: %s' % (os_path, e)) - with atomic_writing(os_path, 'wb') as f: + with atomic_writing(os_path, text=False) as f: f.write(bcontent) def _save_directory(self, os_path, model, name='', path=''): diff --git a/IPython/utils/io.py b/IPython/utils/io.py index 4f624ff..129c2b8 100644 --- a/IPython/utils/io.py +++ b/IPython/utils/io.py @@ -17,6 +17,7 @@ from __future__ import absolute_import #----------------------------------------------------------------------------- import codecs from contextlib import contextmanager +import io import os import sys import tempfile @@ -219,21 +220,55 @@ def temp_pyfile(src, ext='.py'): return fname, f @contextmanager -def atomic_writing(path, mode='w', encoding='utf-8', **kwargs): - tmp_file = path + '.tmp-write' - if 'b' in mode: - encoding = None - - with open(tmp_file, mode, encoding=encoding, **kwargs) as f: - yield f - - # Written successfully, now rename it +def atomic_writing(path, text=True, encoding='utf-8', **kwargs): + """Context manager to write to a file only if the entire write is successful. + + This works by creating a temporary file in the same directory, and renaming + it over the old file if the context is exited without an error. If the + target file is a symlink or a hardlink, this will not be preserved: it will + be replaced by a new regular file. + + On Windows, there is a small chink in the atomicity: the target file is + deleted before renaming the temporary file over it. This appears to be + unavoidable. + + Parameters + ---------- + path : str + The target file to write to. + + text : bool, optional + Whether to open the file in text mode (i.e. to write unicode). Default is + True. + + encoding : str, optional + The encoding to use for files opened in text mode. Default is UTF-8. + + **kwargs + Passed to :func:`io.open`. + """ + dirname, basename = os.path.split(path) + handle, tmp_path = tempfile.mkstemp(prefix=basename, dir=dirname, text=text) + if text: + fileobj = io.open(handle, 'w', encoding=encoding, **kwargs) + else: + fileobj = io.open(handle, 'wb', **kwargs) + + try: + yield fileobj + except: + fileobj.close() + os.remove(tmp_path) + raise + else: + # Written successfully, now rename it + fileobj.close() - if os.name == 'nt' and os.path.exists(path): - # Rename over existing file doesn't work on Windows - os.remove(path) + if os.name == 'nt' and os.path.exists(path): + # Rename over existing file doesn't work on Windows + os.remove(path) - os.rename(tmp_file, path) + os.rename(tmp_path, path) def raw_print(*args, **kw): diff --git a/IPython/utils/tests/test_io.py b/IPython/utils/tests/test_io.py index 9f28ba2..bf4ed7d 100644 --- a/IPython/utils/tests/test_io.py +++ b/IPython/utils/tests/test_io.py @@ -15,6 +15,7 @@ from __future__ import print_function from __future__ import absolute_import import io as stdlib_io +import os.path import sys from subprocess import Popen, PIPE @@ -23,8 +24,11 @@ import unittest import nose.tools as nt from IPython.testing.decorators import skipif -from IPython.utils.io import Tee, capture_output, unicode_std_stream +from IPython.utils.io import (Tee, capture_output, unicode_std_stream, + atomic_writing, + ) from IPython.utils.py3compat import doctest_refactor_print, PY3 +from IPython.utils.tempdir import TemporaryDirectory if PY3: from io import StringIO @@ -121,4 +125,27 @@ def test_UnicodeStdStream_nowrap(): nt.assert_is(unicode_std_stream(), sys.stdout) assert not sys.stdout.closed finally: - sys.stdout = orig_stdout \ No newline at end of file + sys.stdout = orig_stdout + +def test_atomic_writing(): + class CustomExc(Exception): pass + + with TemporaryDirectory() as td: + f1 = os.path.join(td, 'penguin') + with stdlib_io.open(f1, 'w') as f: + f.write(u'Before') + + with nt.assert_raises(CustomExc): + with atomic_writing(f1) as f: + f.write(u'Failing write') + raise CustomExc + + # Because of the exception, the file should not have been modified + with stdlib_io.open(f1, 'r') as f: + nt.assert_equal(f.read(), u'Before') + + with atomic_writing(f1) as f: + f.write(u'Overwritten') + + with stdlib_io.open(f1, 'r') as f: + nt.assert_equal(f.read(), u'Overwritten')