Add atomically_replaced_file for safer output

[bup.git] / lib / bup / helpers.py

diff --git a/lib/bup/helpers.py b/lib/bup/helpers.py
index b55c4a981b4cf462e267b88a5192559c02dc72b5..34772368966ccf633aaf70b22fe47a02368fbee5 100644 (file)
--- a/lib/bup/helpers.py
+++ b/lib/bup/helpers.py
@@ -2,8 +2,9 @@
 
 from ctypes import sizeof, c_void_p
 from os import environ
+from contextlib import contextmanager
 import sys, os, pwd, subprocess, errno, socket, select, mmap, stat, re, struct
-import hashlib, heapq, operator, time, grp
+import hashlib, heapq, operator, time, grp, tempfile
 
 from bup import _helpers
 import bup._helpers as _helpers
@@ -628,6 +629,42 @@ def chunkyreader(f, count = None):
             yield b
 
 
+@contextmanager
+def atomically_replaced_file(name, mode='w', buffering=-1):
+    """Write to file which will atomically become name when finished.
+
+    This contextmanager yields an open file object that is backed by a
+    temporary file which will be renamed (atomically) to the target
+    name if everything succeeds.
+
+    The mode and buffering arguments are handled exactly as with open,
+    and upon success the resulting file will have very restrictive
+    permissions, as per mkstemp.
+
+    E.g.::
+
+        with atomically_replaced_file('foo.txt', 'w') as f:
+            f.write('hello jack.')
+
+    """
+
+    (ffd, tempname) = tempfile.mkstemp(dir=os.path.dirname(name),
+                                       text=('b' not in mode))
+    try:
+        try:
+            f = os.fdopen(ffd, mode, buffering)
+        except:
+            os.close(ffd)
+            raise
+        try:
+            yield f
+        finally:
+            f.close()
+        os.rename(tempname, name)
+    finally:
+        unlink(tempname)  # nonexistant file is ignored
+
+
 def slashappend(s):
     """Append "/" to 's' if it doesn't aleady end in "/"."""
     if s and not s.endswith('/'):