Rework atomic_writing with tests & docstring

IPython/html/services/contents/filemanager.py

@@ -1,532 +1,532 b''
 """A contents manager that uses the local file system for storage."""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import base64
 import io
 import os
 import glob
 import shutil
 
 from tornado import web
 
 from .manager import ContentsManager
 from IPython.nbformat import current
 from IPython.utils.io import atomic_writing
 from IPython.utils.path import ensure_dir_exists
 from IPython.utils.traitlets import Unicode, Bool, TraitError
 from IPython.utils.py3compat import getcwd
 from IPython.utils import tz
 from IPython.html.utils import is_hidden, to_os_path, url_path_join
 
 
 class FileContentsManager(ContentsManager):
 
     root_dir = Unicode(getcwd(), config=True)
 
     save_script = Bool(False, config=True, help='DEPRECATED, IGNORED')
     def _save_script_changed(self):
         self.log.warn("""
         Automatically saving notebooks as scripts has been removed.
         Use `ipython nbconvert --to python [notebook]` instead.
         """)
 
     def _root_dir_changed(self, name, old, new):
         """Do a bit of validation of the root_dir."""
         if not os.path.isabs(new):
             # If we receive a non-absolute path, make it absolute.
             self.root_dir = os.path.abspath(new)
             return
         if not os.path.isdir(new):
             raise TraitError("%r is not a directory" % new)
 
     checkpoint_dir = Unicode('.ipynb_checkpoints', config=True,
         help="""The directory name in which to keep file checkpoints
 
         This is a path relative to the file's own directory.
 
         By default, it is .ipynb_checkpoints
         """
     )
 
     def _copy(self, src, dest):
         """copy src to dest
 
         like shutil.copy2, but log errors in copystat
         """
         shutil.copyfile(src, dest)
         try:
             shutil.copystat(src, dest)
         except OSError as e:
             self.log.debug("copystat on %s failed", dest, exc_info=True)
 
     def _get_os_path(self, name=None, path=''):
         """Given a filename and API path, return its file system
         path.
 
         Parameters
         ----------
         name : string
             A filename
         path : string
             The relative API path to the named file.
 
         Returns
         -------
         path : string
             API path to be evaluated relative to root_dir.
         """
         if name is not None:
             path = url_path_join(path, name)
         return to_os_path(path, self.root_dir)
 
     def path_exists(self, path):
         """Does the API-style path refer to an extant directory?
 
         API-style wrapper for os.path.isdir
 
         Parameters
         ----------
         path : string
             The path to check. This is an API path (`/` separated,
             relative to root_dir).
 
         Returns
         -------
         exists : bool
             Whether the path is indeed a directory.
         """
         path = path.strip('/')
         os_path = self._get_os_path(path=path)
         return os.path.isdir(os_path)
 
     def is_hidden(self, path):
         """Does the API style path correspond to a hidden directory or file?
 
         Parameters
         ----------
         path : string
             The path to check. This is an API path (`/` separated,
             relative to root_dir).
 
         Returns
         -------
         exists : bool
             Whether the path is hidden.
 
         """
         path = path.strip('/')
         os_path = self._get_os_path(path=path)
         return is_hidden(os_path, self.root_dir)
 
     def file_exists(self, name, path=''):
         """Returns True if the file exists, else returns False.
 
         API-style wrapper for os.path.isfile
 
         Parameters
         ----------
         name : string
             The name of the file you are checking.
         path : string
             The relative path to the file's directory (with '/' as separator)
 
         Returns
         -------
         exists : bool
             Whether the file exists.
         """
         path = path.strip('/')
         nbpath = self._get_os_path(name, path=path)
         return os.path.isfile(nbpath)
 
     def exists(self, name=None, path=''):
         """Returns True if the path [and name] exists, else returns False.
 
         API-style wrapper for os.path.exists
 
         Parameters
         ----------
         name : string
             The name of the file you are checking.
         path : string
             The relative path to the file's directory (with '/' as separator)
 
         Returns
         -------
         exists : bool
             Whether the target exists.
         """
         path = path.strip('/')
         os_path = self._get_os_path(name, path=path)
         return os.path.exists(os_path)
 
     def _base_model(self, name, path=''):
         """Build the common base of a contents model"""
         os_path = self._get_os_path(name, path)
         info = os.stat(os_path)
         last_modified = tz.utcfromtimestamp(info.st_mtime)
         created = tz.utcfromtimestamp(info.st_ctime)
         # Create the base model.
         model = {}
         model['name'] = name
         model['path'] = path
         model['last_modified'] = last_modified
         model['created'] = created
         model['content'] = None
         model['format'] = None
         return model
 
     def _dir_model(self, name, path='', content=True):
         """Build a model for a directory
 
         if content is requested, will include a listing of the directory
         """
         os_path = self._get_os_path(name, path)
 
         four_o_four = u'directory does not exist: %r' % os_path
 
         if not os.path.isdir(os_path):
             raise web.HTTPError(404, four_o_four)
         elif is_hidden(os_path, self.root_dir):
             self.log.info("Refusing to serve hidden directory %r, via 404 Error",
                 os_path
             )
             raise web.HTTPError(404, four_o_four)
 
         if name is None:
             if '/' in path:
                 path, name = path.rsplit('/', 1)
             else:
                 name = ''
         model = self._base_model(name, path)
         model['type'] = 'directory'
         dir_path = u'{}/{}'.format(path, name)
         if content:
             model['content'] = contents = []
             for os_path in glob.glob(self._get_os_path('*', dir_path)):
                 name = os.path.basename(os_path)
                 if self.should_list(name) and not is_hidden(os_path, self.root_dir):
                     contents.append(self.get_model(name=name, path=dir_path, content=False))
 
             model['format'] = 'json'
 
         return model
 
     def _file_model(self, name, path='', content=True):
         """Build a model for a file
 
         if content is requested, include the file contents.
         UTF-8 text files will be unicode, binary files will be base64-encoded.
         """
         model = self._base_model(name, path)
         model['type'] = 'file'
         if content:
             os_path = self._get_os_path(name, path)
             with io.open(os_path, 'rb') as f:
                 bcontent = f.read()
             try:
                 model['content'] = bcontent.decode('utf8')
             except UnicodeError as e:
                 model['content'] = base64.encodestring(bcontent).decode('ascii')
                 model['format'] = 'base64'
             else:
                 model['format'] = 'text'
         return model
 
 
     def _notebook_model(self, name, path='', content=True):
         """Build a notebook model
 
         if content is requested, the notebook content will be populated
         as a JSON structure (not double-serialized)
         """
         model = self._base_model(name, path)
         model['type'] = 'notebook'
         if content:
             os_path = self._get_os_path(name, path)
             with io.open(os_path, 'r', encoding='utf-8') as f:
                 try:
                     nb = current.read(f, u'json')
                 except Exception as e:
                     raise web.HTTPError(400, u"Unreadable Notebook: %s %s" % (os_path, e))
             self.mark_trusted_cells(nb, name, path)
             model['content'] = nb
             model['format'] = 'json'
         return model
 
     def get_model(self, name, path='', content=True):
         """ Takes a path and name for an entity and returns its model
 
         Parameters
         ----------
         name : str
             the name of the target
         path : str
             the API path that describes the relative path for the target
 
         Returns
         -------
         model : dict
             the contents model. If content=True, returns the contents
             of the file or directory as well.
         """
         path = path.strip('/')
 
         if not self.exists(name=name, path=path):
             raise web.HTTPError(404, u'No such file or directory: %s/%s' % (path, name))
 
         os_path = self._get_os_path(name, path)
         if os.path.isdir(os_path):
             model = self._dir_model(name, path, content)
         elif name.endswith('.ipynb'):
             model = self._notebook_model(name, path, content)
         else:
             model = self._file_model(name, path, content)
         return model
 
     def _save_notebook(self, os_path, model, name='', path=''):
         """save a notebook file"""
         # Save the notebook file
         nb = current.to_notebook_json(model['content'])
 
         self.check_and_sign(nb, name, path)
 
         if 'name' in nb['metadata']:
             nb['metadata']['name'] = u''
 
         with atomic_writing(os_path, encoding='utf-8') as f:
             current.write(nb, f, u'json')
 
     def _save_file(self, os_path, model, name='', path=''):
         """save a non-notebook file"""
         fmt = model.get('format', None)
         if fmt not in {'text', 'base64'}:
             raise web.HTTPError(400, "Must specify format of file contents as 'text' or 'base64'")
         try:
             content = model['content']
             if fmt == 'text':
                 bcontent = content.encode('utf8')
             else:
                 b64_bytes = content.encode('ascii')
                 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=''):
         """create a directory"""
         if is_hidden(os_path, self.root_dir):
             raise web.HTTPError(400, u'Cannot create hidden directory %r' % os_path)
         if not os.path.exists(os_path):
             os.mkdir(os_path)
         elif not os.path.isdir(os_path):
             raise web.HTTPError(400, u'Not a directory: %s' % (os_path))
         else:
             self.log.debug("Directory %r already exists", os_path)
 
     def save(self, model, name='', path=''):
         """Save the file model and return the model with no content."""
         path = path.strip('/')
 
         if 'type' not in model:
             raise web.HTTPError(400, u'No file type provided')
         if 'content' not in model and model['type'] != 'directory':
             raise web.HTTPError(400, u'No file content provided')
 
         # One checkpoint should always exist
         if self.file_exists(name, path) and not self.list_checkpoints(name, path):
             self.create_checkpoint(name, path)
 
         new_path = model.get('path', path).strip('/')
         new_name = model.get('name', name)
 
         if path != new_path or name != new_name:
             self.rename(name, path, new_name, new_path)
 
         os_path = self._get_os_path(new_name, new_path)
         self.log.debug("Saving %s", os_path)
         try:
             if model['type'] == 'notebook':
                 self._save_notebook(os_path, model, new_name, new_path)
             elif model['type'] == 'file':
                 self._save_file(os_path, model, new_name, new_path)
             elif model['type'] == 'directory':
                 self._save_directory(os_path, model, new_name, new_path)
             else:
                 raise web.HTTPError(400, "Unhandled contents type: %s" % model['type'])
         except web.HTTPError:
             raise
         except Exception as e:
             raise web.HTTPError(400, u'Unexpected error while saving file: %s %s' % (os_path, e))
 
         model = self.get_model(new_name, new_path, content=False)
         return model
 
     def update(self, model, name, path=''):
         """Update the file's path and/or name
 
         For use in PATCH requests, to enable renaming a file without
         re-uploading its contents. Only used for renaming at the moment.
         """
         path = path.strip('/')
         new_name = model.get('name', name)
         new_path = model.get('path', path).strip('/')
         if path != new_path or name != new_name:
             self.rename(name, path, new_name, new_path)
         model = self.get_model(new_name, new_path, content=False)
         return model
 
     def delete(self, name, path=''):
         """Delete file by name and path."""
         path = path.strip('/')
         os_path = self._get_os_path(name, path)
         rm = os.unlink
         if os.path.isdir(os_path):
             listing = os.listdir(os_path)
             # don't delete non-empty directories (checkpoints dir doesn't count)
             if listing and listing != [self.checkpoint_dir]:
                 raise web.HTTPError(400, u'Directory %s not empty' % os_path)
         elif not os.path.isfile(os_path):
             raise web.HTTPError(404, u'File does not exist: %s' % os_path)
 
         # clear checkpoints
         for checkpoint in self.list_checkpoints(name, path):
             checkpoint_id = checkpoint['id']
             cp_path = self.get_checkpoint_path(checkpoint_id, name, path)
             if os.path.isfile(cp_path):
                 self.log.debug("Unlinking checkpoint %s", cp_path)
                 os.unlink(cp_path)
 
         if os.path.isdir(os_path):
             self.log.debug("Removing directory %s", os_path)
             shutil.rmtree(os_path)
         else:
             self.log.debug("Unlinking file %s", os_path)
             rm(os_path)
 
     def rename(self, old_name, old_path, new_name, new_path):
         """Rename a file."""
         old_path = old_path.strip('/')
         new_path = new_path.strip('/')
         if new_name == old_name and new_path == old_path:
             return
 
         new_os_path = self._get_os_path(new_name, new_path)
         old_os_path = self._get_os_path(old_name, old_path)
 
         # Should we proceed with the move?
         if os.path.isfile(new_os_path):
             raise web.HTTPError(409, u'File with name already exists: %s' % new_os_path)
 
         # Move the file
         try:
             shutil.move(old_os_path, new_os_path)
         except Exception as e:
             raise web.HTTPError(500, u'Unknown error renaming file: %s %s' % (old_os_path, e))
 
         # Move the checkpoints
         old_checkpoints = self.list_checkpoints(old_name, old_path)
         for cp in old_checkpoints:
             checkpoint_id = cp['id']
             old_cp_path = self.get_checkpoint_path(checkpoint_id, old_name, old_path)
             new_cp_path = self.get_checkpoint_path(checkpoint_id, new_name, new_path)
             if os.path.isfile(old_cp_path):
                 self.log.debug("Renaming checkpoint %s -> %s", old_cp_path, new_cp_path)
                 shutil.move(old_cp_path, new_cp_path)
 
     # Checkpoint-related utilities
 
     def get_checkpoint_path(self, checkpoint_id, name, path=''):
         """find the path to a checkpoint"""
         path = path.strip('/')
         basename, ext = os.path.splitext(name)
         filename = u"{name}-{checkpoint_id}{ext}".format(
             name=basename,
             checkpoint_id=checkpoint_id,
             ext=ext,
         )
         os_path = self._get_os_path(path=path)
         cp_dir = os.path.join(os_path, self.checkpoint_dir)
         ensure_dir_exists(cp_dir)
         cp_path = os.path.join(cp_dir, filename)
         return cp_path
 
     def get_checkpoint_model(self, checkpoint_id, name, path=''):
         """construct the info dict for a given checkpoint"""
         path = path.strip('/')
         cp_path = self.get_checkpoint_path(checkpoint_id, name, path)
         stats = os.stat(cp_path)
         last_modified = tz.utcfromtimestamp(stats.st_mtime)
         info = dict(
             id = checkpoint_id,
             last_modified = last_modified,
         )
         return info
 
     # public checkpoint API
 
     def create_checkpoint(self, name, path=''):
         """Create a checkpoint from the current state of a file"""
         path = path.strip('/')
         src_path = self._get_os_path(name, path)
         # only the one checkpoint ID:
         checkpoint_id = u"checkpoint"
         cp_path = self.get_checkpoint_path(checkpoint_id, name, path)
         self.log.debug("creating checkpoint for %s", name)
         self._copy(src_path, cp_path)
 
         # return the checkpoint info
         return self.get_checkpoint_model(checkpoint_id, name, path)
 
     def list_checkpoints(self, name, path=''):
         """list the checkpoints for a given file
 
         This contents manager currently only supports one checkpoint per file.
         """
         path = path.strip('/')
         checkpoint_id = "checkpoint"
         os_path = self.get_checkpoint_path(checkpoint_id, name, path)
         if not os.path.exists(os_path):
             return []
         else:
             return [self.get_checkpoint_model(checkpoint_id, name, path)]
 
 
     def restore_checkpoint(self, checkpoint_id, name, path=''):
         """restore a file to a checkpointed state"""
         path = path.strip('/')
         self.log.info("restoring %s from checkpoint %s", name, checkpoint_id)
         nb_path = self._get_os_path(name, path)
         cp_path = self.get_checkpoint_path(checkpoint_id, name, path)
         if not os.path.isfile(cp_path):
             self.log.debug("checkpoint file does not exist: %s", cp_path)
             raise web.HTTPError(404,
                 u'checkpoint does not exist: %s-%s' % (name, checkpoint_id)
             )
         # ensure notebook is readable (never restore from an unreadable notebook)
         if cp_path.endswith('.ipynb'):
             with io.open(cp_path, 'r', encoding='utf-8') as f:
                 current.read(f, u'json')
         self._copy(cp_path, nb_path)
         self.log.debug("copying %s -> %s", cp_path, nb_path)
 
     def delete_checkpoint(self, checkpoint_id, name, path=''):
         """delete a file's checkpoint"""
         path = path.strip('/')
         cp_path = self.get_checkpoint_path(checkpoint_id, name, path)
         if not os.path.isfile(cp_path):
             raise web.HTTPError(404,
                 u'Checkpoint does not exist: %s%s-%s' % (path, name, checkpoint_id)
             )
         self.log.debug("unlinking %s", cp_path)
         os.unlink(cp_path)
 
     def info_string(self):
         return "Serving notebooks from local directory: %s" % self.root_dir
 
     def get_kernel_path(self, name, path='', model=None):
         """Return the initial working dir a kernel associated with a given notebook"""
         return os.path.join(self.root_dir, path)

IPython/utils/io.py

@@ -1,280 +1,315 b''
 # encoding: utf-8
 """
 IO related utilities.
 """
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2008-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 from __future__ import print_function
 from __future__ import absolute_import
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 import codecs
 from contextlib import contextmanager
+import io
 import os
 import sys
 import tempfile
 from .capture import CapturedIO, capture_output
 from .py3compat import string_types, input, PY3
 
 #-----------------------------------------------------------------------------
 # Code
 #-----------------------------------------------------------------------------
 
 
 class IOStream:
 
     def __init__(self,stream, fallback=None):
         if not hasattr(stream,'write') or not hasattr(stream,'flush'):
             if fallback is not None:
                 stream = fallback
             else:
                 raise ValueError("fallback required, but not specified")
         self.stream = stream
         self._swrite = stream.write
 
         # clone all methods not overridden:
         def clone(meth):
             return not hasattr(self, meth) and not meth.startswith('_')
         for meth in filter(clone, dir(stream)):
             setattr(self, meth, getattr(stream, meth))
 
     def __repr__(self):
         cls = self.__class__
         tpl = '{mod}.{cls}({args})'
         return tpl.format(mod=cls.__module__, cls=cls.__name__, args=self.stream)
 
     def write(self,data):
         try:
             self._swrite(data)
         except:
             try:
                 # print handles some unicode issues which may trip a plain
                 # write() call.  Emulate write() by using an empty end
                 # argument.
                 print(data, end='', file=self.stream)
             except:
                 # if we get here, something is seriously broken.
                 print('ERROR - failed to write data to stream:', self.stream,
                       file=sys.stderr)
 
     def writelines(self, lines):
         if isinstance(lines, string_types):
             lines = [lines]
         for line in lines:
             self.write(line)
 
     # This class used to have a writeln method, but regular files and streams
     # in Python don't have this method. We need to keep this completely
     # compatible so we removed it.
 
     @property
     def closed(self):
         return self.stream.closed
 
     def close(self):
         pass
 
 # setup stdin/stdout/stderr to sys.stdin/sys.stdout/sys.stderr
 devnull = open(os.devnull, 'w') 
 stdin = IOStream(sys.stdin, fallback=devnull)
 stdout = IOStream(sys.stdout, fallback=devnull)
 stderr = IOStream(sys.stderr, fallback=devnull)
 
 class IOTerm:
     """ Term holds the file or file-like objects for handling I/O operations.
 
     These are normally just sys.stdin, sys.stdout and sys.stderr but for
     Windows they can can replaced to allow editing the strings before they are
     displayed."""
 
     # In the future, having IPython channel all its I/O operations through
     # this class will make it easier to embed it into other environments which
     # are not a normal terminal (such as a GUI-based shell)
     def __init__(self, stdin=None, stdout=None, stderr=None):
         mymodule = sys.modules[__name__]
         self.stdin  = IOStream(stdin, mymodule.stdin)
         self.stdout = IOStream(stdout, mymodule.stdout)
         self.stderr = IOStream(stderr, mymodule.stderr)
 
 
 class Tee(object):
     """A class to duplicate an output stream to stdout/err.
 
     This works in a manner very similar to the Unix 'tee' command.
 
     When the object is closed or deleted, it closes the original file given to
     it for duplication.
     """
     # Inspired by:
     # http://mail.python.org/pipermail/python-list/2007-May/442737.html
 
     def __init__(self, file_or_name, mode="w", channel='stdout'):
         """Construct a new Tee object.
 
         Parameters
         ----------
         file_or_name : filename or open filehandle (writable)
           File that will be duplicated
 
         mode : optional, valid mode for open().
           If a filename was give, open with this mode.
 
         channel : str, one of ['stdout', 'stderr']
         """
         if channel not in ['stdout', 'stderr']:
             raise ValueError('Invalid channel spec %s' % channel)
 
         if hasattr(file_or_name, 'write') and hasattr(file_or_name, 'seek'):
             self.file = file_or_name
         else:
             self.file = open(file_or_name, mode)
         self.channel = channel
         self.ostream = getattr(sys, channel)
         setattr(sys, channel, self)
         self._closed = False
 
     def close(self):
         """Close the file and restore the channel."""
         self.flush()
         setattr(sys, self.channel, self.ostream)
         self.file.close()
         self._closed = True
 
     def write(self, data):
         """Write data to both channels."""
         self.file.write(data)
         self.ostream.write(data)
         self.ostream.flush()
 
     def flush(self):
         """Flush both channels."""
         self.file.flush()
         self.ostream.flush()
 
     def __del__(self):
         if not self._closed:
             self.close()
 
 
 def ask_yes_no(prompt, default=None, interrupt=None):
     """Asks a question and returns a boolean (y/n) answer.
 
     If default is given (one of 'y','n'), it is used if the user input is
     empty. If interrupt is given (one of 'y','n'), it is used if the user
     presses Ctrl-C. Otherwise the question is repeated until an answer is
     given.
 
     An EOF is treated as the default answer.  If there is no default, an
     exception is raised to prevent infinite loops.
 
     Valid answers are: y/yes/n/no (match is not case sensitive)."""
 
     answers = {'y':True,'n':False,'yes':True,'no':False}
     ans = None
     while ans not in answers.keys():
         try:
             ans = input(prompt+' ').lower()
             if not ans:  # response was an empty string
                 ans = default
         except KeyboardInterrupt:
             if interrupt:
                 ans = interrupt
         except EOFError:
             if default in answers.keys():
                 ans = default
                 print()
             else:
                 raise
 
     return answers[ans]
 
 
 def temp_pyfile(src, ext='.py'):
     """Make a temporary python file, return filename and filehandle.
 
     Parameters
     ----------
     src : string or list of strings (no need for ending newlines if list)
       Source code to be written to the file.
 
     ext : optional, string
       Extension for the generated file.
 
     Returns
     -------
     (filename, open filehandle)
       It is the caller's responsibility to close the open file and unlink it.
     """
     fname = tempfile.mkstemp(ext)[1]
     f = open(fname,'w')
     f.write(src)
     f.flush()
     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):
     """Raw print to sys.__stdout__, otherwise identical interface to print()."""
 
     print(*args, sep=kw.get('sep', ' '), end=kw.get('end', '\n'),
           file=sys.__stdout__)
     sys.__stdout__.flush()
 
 
 def raw_print_err(*args, **kw):
     """Raw print to sys.__stderr__, otherwise identical interface to print()."""
 
     print(*args, sep=kw.get('sep', ' '), end=kw.get('end', '\n'),
           file=sys.__stderr__)
     sys.__stderr__.flush()
 
 
 # Short aliases for quick debugging, do NOT use these in production code.
 rprint = raw_print
 rprinte = raw_print_err
 
 def unicode_std_stream(stream='stdout'):
     u"""Get a wrapper to write unicode to stdout/stderr as UTF-8.
 
     This ignores environment variables and default encodings, to reliably write
     unicode to stdout or stderr.
 
     ::
 
         unicode_std_stream().write(u'ł@e¶ŧ←')
     """
     assert stream in ('stdout', 'stderr')
     stream  = getattr(sys, stream)
     if PY3:
         try:
             stream_b = stream.buffer
         except AttributeError:
             # sys.stdout has been replaced - use it directly
             return stream
     else:
         stream_b = stream
 
     return codecs.getwriter('utf-8')(stream_b)

IPython/utils/tests/test_io.py

@@ -1,124 +1,151 b''
 # encoding: utf-8
 """Tests for io.py"""
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2008-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 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
 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
 else:
     from StringIO import StringIO
 
 #-----------------------------------------------------------------------------
 # Tests
 #-----------------------------------------------------------------------------
 
 
 def test_tee_simple():
     "Very simple check with stdout only"
     chan = StringIO()
     text = 'Hello'
     tee = Tee(chan, channel='stdout')
     print(text, file=chan)
     nt.assert_equal(chan.getvalue(), text+"\n")
 
 
 class TeeTestCase(unittest.TestCase):
 
     def tchan(self, channel, check='close'):
         trap = StringIO()
         chan = StringIO()
         text = 'Hello'
         
         std_ori = getattr(sys, channel)
         setattr(sys, channel, trap)
 
         tee = Tee(chan, channel=channel)
         print(text, end='', file=chan)
         setattr(sys, channel, std_ori)
         trap_val = trap.getvalue()
         nt.assert_equal(chan.getvalue(), text)
         if check=='close':
             tee.close()
         else:
             del tee
 
     def test(self):
         for chan in ['stdout', 'stderr']:
             for check in ['close', 'del']:
                 self.tchan(chan, check)
 
 def test_io_init():
     """Test that io.stdin/out/err exist at startup"""
     for name in ('stdin', 'stdout', 'stderr'):
         cmd = doctest_refactor_print("from IPython.utils import io;print io.%s.__class__"%name)
         p = Popen([sys.executable, '-c', cmd],
                     stdout=PIPE)
         p.wait()
         classname = p.stdout.read().strip().decode('ascii')
         # __class__ is a reference to the class object in Python 3, so we can't
         # just test for string equality.
         assert 'IPython.utils.io.IOStream' in classname, classname
 
 def test_capture_output():
     """capture_output() context works"""
     
     with capture_output() as io:
         print('hi, stdout')
         print('hi, stderr', file=sys.stderr)
     
     nt.assert_equal(io.stdout, 'hi, stdout\n')
     nt.assert_equal(io.stderr, 'hi, stderr\n')
 
 def test_UnicodeStdStream():
     # Test wrapping a bytes-level stdout
     if PY3:
         stdoutb = stdlib_io.BytesIO()
         stdout = stdlib_io.TextIOWrapper(stdoutb, encoding='ascii')
     else:
         stdout = stdoutb = stdlib_io.BytesIO()
 
     orig_stdout = sys.stdout
     sys.stdout = stdout
     try:
         sample = u"@łe¶ŧ←"
         unicode_std_stream().write(sample)
 
         output = stdoutb.getvalue().decode('utf-8')
         nt.assert_equal(output, sample)
         assert not stdout.closed
     finally:
         sys.stdout = orig_stdout
 
 @skipif(not PY3, "Not applicable on Python 2")
 def test_UnicodeStdStream_nowrap():
     # If we replace stdout with a StringIO, it shouldn't get wrapped.
     orig_stdout = sys.stdout
     sys.stdout = StringIO()
     try:
         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')