Use a modification context object.

1 files changed, 44 insertions, 29 deletions

diff --git a/deheader b/deheader
index 7ab4028..ee1fa49 100755
--- a/deheader
+++ b/deheader
@@ -112,6 +112,27 @@ class InclusionMap:
         "Undo forgetting of a dependency."
         self.c_to_h[sourcefile].append(header)
 
+class SaveForModification:
+    "Prepare a file to be temporarily modified, with guaranteed reversion."
+    def __init__(self, filename):
+        self.filename = filename
+        self.original = filename + ".orig"
+        os.rename(self.filename, self.original)
+    def remove_header(self, header):
+        "Prepare a version with a specified header deleted."
+        ofp = open(self.filename, "w")
+        for line in open(self.original):
+            if line != header:
+                ofp.write(line)
+        ofp.close()
+    def __del__(self):
+        "Revert modifications on the file at the end of this onject's lifetime."
+        try:
+            os.remove(self.filename)
+        except OSError:
+            pass
+        os.rename(self.original, self.filename)
+
 def trim(line):
     "Get file reference from an #include, retaining <> if a system header."
     trimmed = line[9:].strip()
@@ -145,44 +166,19 @@ def compile(source, msg="", verbosity=0):
           % (sourcefile, msg, explain, end-start)
     return (status, end - start)
 
-def deheader(sourcefile, includes, remove, verbose):
-    # Sanity check against broken sourcefiles; we want this to
-    # complain visibly if the sourcefile won't build at all.
-    (st, t) = compile(sourcefile, min(1, verbose))
-    if st == 0:
-        # Now do the analysis
-        if sourcefile.endswith(".c") or sourcefile.endswith(".cpp"):
-            unneeded = c_analyze(sourcefile, includes, verbose)
-        if unneeded:
-            for line in unneeded:
-                print "deheader: remove %s from %s" % (trim(line), sourcefile)
-            if remove:
-                os.rename(sourcefile, original)
-                for header in includes:
-                    ofp = open(sourcefile, "w")
-                    for line in open(original):
-                        if line not in unneeded:
-                            ofp.write(line)
-                    ofp.close()
-
 def c_analyze(source, includes, verbosity):
     "Given a C file and a list of includes, return those that can be omitted."
     # We'll remove headers in reverse order, because later unnecessary
     # headers might depend on earlier ones
     includes.reverse()
     unneeded = []
-    original = sourcefile + ".orig"
     baton = Baton(sourcefile + ": ", "Done")
-    os.rename(sourcefile, original)
+    saveit = SaveForModification(sourcefile)
     try:
         while True:                
             keepgoing = False
             for header in includes:
-                ofp = open(sourcefile, "w")
-                for line in open(original):
-                    if line != header:
-                        ofp.write(line)
-                ofp.close()
+                saveit.remove_header(header)
                 (st, t) = compile(sourcefile, " without %s" % trim(header), verbosity)
                 if st == 0:
                     unneeded.append(header)
@@ -193,11 +189,30 @@ def c_analyze(source, includes, verbosity):
             if not keepgoing:
                 break
     finally:
+        del saveit
         baton.end()
-        os.remove(sourcefile)
-        os.rename(original, sourcefile)
     return unneeded
 
+def deheader(sourcefile, includes, remove, verbose):
+    # Sanity check against broken sourcefiles; we want this to
+    # complain visibly if the sourcefile won't build at all.
+    (st, t) = compile(sourcefile, min(1, verbose))
+    if st == 0:
+        # Now do the analysis
+        if sourcefile.endswith(".c") or sourcefile.endswith(".cpp"):
+            unneeded = c_analyze(sourcefile, includes, verbose)
+        if unneeded:
+            for line in unneeded:
+                print "deheader: remove %s from %s" % (trim(line), sourcefile)
+            if remove:
+                os.rename(sourcefile, original)
+                for header in includes:
+                    ofp = open(sourcefile, "w")
+                    for line in open(original):
+                        if line not in unneeded:
+                            ofp.write(line)
+                    ofp.close()
+
 if __name__ == "__main__":
     (options, arguments) = getopt.getopt(sys.argv[1:], "hrvi:",
                                          ["help", "ignore",