aopen.py

import os.path
import tempfile
import StringIO
import os
import sys


class AtomicOpen(StringIO.StringIO, object):
    """A file-like context manager for atomicly writing files
    """
    def __init__(self, filename, *p, **k):
        self._realname = os.path.abspath(filename)
        self._dir = os.path.dirname(self._realname)
        super(AtomicOpen, self).__init__(*p, **k)

    def __enter__(self):
        try:
            with open(self._realname) as f:
                for x in f.read():
                    self.write(x)
            self.seek(0)
        except IOError, e:
            if e[0] == 2:
                pass
            else:
                raise
        return self

    def __exit__(self, exp, value, traceback):
        with tempfile.NamedTemporaryFile(dir=self._dir, delete=False) as tmp:
            tmp.write(self.getvalue())
            name = tmp.name
            tmp.file.flush()
            os.fsync(tmp.file.fileno())
        fd = os.open(self._dir, 0)
        os.fsync(fd)
        os.close(fd)
        if sys.platform == 'win32':
            try:
                os.unlink(self._realname)
            except IOError, e:
                if e[0] == 2:
                    pass
                else:
                    raise
        os.rename(name, self._realname)
        StringIO.StringIO.close(self)