scripts: modernize logextract.py (with opens, generators)

Signed-off-by: Maël Kerbiriou <m431kerbiriou@gmail.com>

1 files changed, 37 insertions, 35 deletions

diff --git a/scripts/logextract.py b/scripts/logextract.py
index 96a924d6..592f6f8a 100755
--- a/scripts/logextract.py
+++ b/scripts/logextract.py
@@ -4,7 +4,7 @@
 # Copyright (C) 2017  Kevin O'Connor <kevin@koconnor.net>
 #
 # This file may be distributed under the terms of the GNU GPLv3 license.
-import sys, re, collections, ast
+import sys, re, collections, ast, itertools
 
 def format_comment(line_num, line):
     return "# %6d: %s" % (line_num, line)
@@ -40,9 +40,10 @@ class GatherConfig:
         if comment is not None:
             self.comments.append(comment)
     def write_file(self):
-        f = open(self.filename, 'wt')
-        f.write('\n'.join(self.comments + self.config_lines).strip() + '\n')
-        f.close()
+        lines = itertools.chain(self.comments, self.config_lines)
+        lines = ('%s\n' % l for l in lines)
+        with open(self.filename, 'wt') as f:
+            f.writelines(lines)
 
 
 ######################################################################
@@ -383,10 +384,9 @@ class GCodeStream:
     def get_lines(self):
         # Produce output gcode stream
         if self.gcode_stream:
-            data = [ast.literal_eval(gc) for gc in self.gcode_commands]
-            f = open(self.gcode_filename, 'wt')
-            f.write(self.gcode_state + ''.join(data))
-            f.close()
+            data = (ast.literal_eval(gc) for gc in self.gcode_commands)
+            with open(self.gcode_filename, 'wt') as f:
+                f.write(self.gcode_state + ''.join(data))
         return self.gcode_stream
 
 api_cmd_r = re.compile(r"^Received " + time_s + r": \{.*\}$")
@@ -566,10 +566,10 @@ class GatherShutdown:
         # Produce output sorted by timestamp
         out = [i for s in streams for i in s]
         out.sort()
-        out = [i[2] for i in out]
-        f = open(self.filename, 'wt')
-        f.write('\n'.join(self.comments + out))
-        f.close()
+        lines = itertools.chain(self.comments, (i[2] for i in out))
+        lines = ('%s\n' % l for l in lines)
+        with open(self.filename, 'wt') as f:
+            f.writelines(lines)
 
 
 ######################################################################
@@ -583,29 +583,31 @@ def main():
     handler = None
     recent_lines = collections.deque([], 200)
     # Parse log file
-    f = open(logname, 'rt')
-    for line_num, line in enumerate(f):
-        line = line.rstrip()
-        line_num += 1
-        recent_lines.append((line_num, line))
-        if handler is not None:
-            ret = handler.add_line(line_num, line)
-            if ret:
-                continue
-            recent_lines.clear()
-            handler = None
-        if line.startswith('Git version'):
-            last_git = format_comment(line_num, line)
-        elif line.startswith('Start printer at'):
-            last_start = format_comment(line_num, line)
-        elif line == '===== Config file =====':
-            handler = GatherConfig(configs, line_num, recent_lines, logname)
-            handler.add_comment(last_git)
-            handler.add_comment(last_start)
-        elif 'shutdown: ' in line or line.startswith('Dumping '):
-            handler = GatherShutdown(configs, line_num, recent_lines, logname)
-            handler.add_comment(last_git)
-            handler.add_comment(last_start)
+    with open(logname, 'rt') as f:
+        for line_num, line in enumerate(f):
+            line = line.rstrip()
+            line_num += 1
+            recent_lines.append((line_num, line))
+            if handler is not None:
+                ret = handler.add_line(line_num, line)
+                if ret:
+                    continue
+                recent_lines.clear()
+                handler = None
+            if line.startswith('Git version'):
+                last_git = format_comment(line_num, line)
+            elif line.startswith('Start printer at'):
+                last_start = format_comment(line_num, line)
+            elif line == '===== Config file =====':
+                handler = GatherConfig(configs, line_num,
+                                       recent_lines, logname)
+                handler.add_comment(last_git)
+                handler.add_comment(last_start)
+            elif 'shutdown: ' in line or line.startswith('Dumping '):
+                handler = GatherShutdown(configs, line_num,
+                                         recent_lines, logname)
+                handler.add_comment(last_git)
+                handler.add_comment(last_start)
     if handler is not None:
         handler.finalize()
     # Write found config files