Added utils.AtomicFile and converted our uses of 'w' to it. This rocks.

plugins/Infobot.py

@@ -104,7 +104,7 @@ class InfobotDB(object):
                           'I hear ya']
 
     def flush(self):
-        fd = file(filename, 'w')
+        fd = utils.AtomicFile(filename)
         pickle.dump((self._is, self._are), fd)
         fd.close()
 

plugins/URL.py

@@ -149,8 +149,7 @@ class URLDB(object):
         return [url for (url, nick) in self.getUrlsAndNicks(p)]
 
     def vacuum(self):
-        filename = utils.mktemp()
-        out = file(filename, 'w')
+        out = utils.AtomicFile(self.filename)
         notAdded = 0
         urls = self.getUrlsAndNicks(lambda *args: True)
         seen = sets.Set()
@@ -165,7 +164,6 @@ class URLDB(object):
             if urlNick is not None:
                 out.write(self._formatRecord(*urlNick))
         out.close()
-        shutil.move(filename, self.filename)
         self.log.info('Vacuumed %s, removed %s records.',
                       self.filename, notAdded)
                 

plugins/__init__.py

@@ -221,9 +221,14 @@ class ChannelUserDB(ChannelUserDictionary):
             log.debug('Exception: %s', utils.exnToString(e))
 
     def flush(self):
-        fd = file(self.filename, 'w')
+        fd = utils.AtomicFile(self.filename)
         writer = csv.writer(fd)
         items = self.items()
+        if not items:
+            log.warning('%s: Refusing to write blank file.',
+                        self.__class__.__name__)
+            fd.rollback()
+            return
         items.sort()
         for ((channel, id), v) in items:
             L = self.serialize(v)

src/cdb.py

@@ -136,7 +136,7 @@ def make(dbFilename, readFilename=None):
 class Maker(object):
     """Class for making CDB databases."""
     def __init__(self, filename):
-        self.fd = file(filename, 'w')
+        self.fd = utils.AtomicFile(filename)
         self.filename = filename
         self.fd.seek(2048)
         self.hashPointers = [(0, 0)] * 256

src/ircdb.py

@@ -571,7 +571,7 @@ class UsersDictionary(utils.IterableMap):
         if self.filename is not None:
             L = self.users.items()
             L.sort()
-            fd = file(self.filename, 'w')
+            fd = utils.AtomicFile(self.filename)
             for (id, u) in L:
                 fd.write('user %s' % id)
                 fd.write(os.linesep)
@@ -732,7 +732,7 @@ class ChannelsDictionary(utils.IterableMap):
     def flush(self):
         """Flushes the channel database to its file."""
         if self.filename is not None:
-            fd = file(self.filename, 'w')
+            fd = utils.AtomicFile(self.filename)
             for (channel, c) in self.channels.iteritems():
                 fd.write('channel %s' % channel)
                 fd.write(os.linesep)
@@ -792,7 +792,7 @@ class IgnoresDB(object):
 
     def flush(self):
         if self.filename is not None:
-            fd = file(self.filename, 'w')
+            fd = utils.AtomicFile(self.filename)
             for hostmask in self.hostmasks:
                 fd.write(hostmask)
                 fd.write(os.linesep)

src/ircutils.py

@@ -472,7 +472,6 @@ class IrcString(str):
         x.lowered = toLower(x)
         return x
 
-
     def __eq__(self, s):
         try:
             return toLower(s) == self.lowered

src/registry.py

@@ -78,7 +78,7 @@ def open(filename, clear=False):
 def close(registry, filename, annotated=True, helpOnceOnly=False):
     first = True
     helpCache = sets.Set()
-    fd = file(filename, 'w')
+    fd = utils.AtomicFile(filename)
     for (name, value) in registry.getValues(getChildren=True):
         if annotated and hasattr(value,'help') and value.help:
             if not helpOnceOnly or value.help not in self.helpCache:

src/utils.py

@@ -676,18 +676,38 @@ def stackTrace():
 
 class AtomicFile(file):
     """Used for files that need to be atomically written -- i.e., if there's a
-    failure, the original file remains, unmodified."""
-    def __init__(self, filename, flags='w'):
-        if flags not in ('a', 'w'):
-            raise ValueError, 'AtomicFile should only be used for writing.'
+    failure, the original file remains, unmodified.
+
+    Opens the file in 'w' mode."""
+    def __init__(self, filename, allowEmptyOverwrite=False):
         self.filename = filename
+        self.rolledback = False
+        self.allowEmptyOverwrite = allowEmptyOverwrite
         self.tempFilename = '%s.%s' % (filename, mktemp())
-        super(AtomicFile, self).__init__(self.tempFilename, flags)
+        super(AtomicFile, self).__init__(self.tempFilename, 'w')
+
+    def rollback(self):
+        #print 'AtomicFile.rollback'
+        super(AtomicFile, self).close()
+        if os.path.exists(self.tempFilename):
+            print 'AtomicFile: Removing %s.' % self.tempFilename
+            os.remove(self.tempFilename)
+        self.rolledback = True
 
     def close(self):
+        #print 'AtomicFile.close'
+        if not self.rolledback:
+            #print 'AtomicFile.close: actually closing.'
             super(AtomicFile, self).close()
+            size = os.stat(self.tempFilename).st_size
+            if size or self.allowEmptyOverwrite:
+                if os.path.exists(self.tempFilename):
                     shutil.move(self.tempFilename, self.filename)
 
+    def __del__(self):
+        #print 'AtomicFile.__del__'
+        self.rollback()
+
 if __name__ == '__main__':
     import sys, doctest
     doctest.testmod(sys.modules['__main__'])