Make stitch_net_log_files.py handle more cases.

 * Order event files from oldest to newest
 * Remove the --num_files switch
 * Don't fail if end_netlog.json is missing (since it probably will be
   when using this)

Bug: 744815
Change-Id: I26307d3c293eb445b7e3456d6b359f8ef05a0204
Reviewed-on: https://chromium-review.googlesource.com/578484
Commit-Queue: Eric Roman <eroman@chromium.org>
Reviewed-by: Helen Li <xunjieli@chromium.org>
Cr-Commit-Position: refs/heads/master@{#488045}

net/tools/stitch_net_log_files.py

@@ -4,50 +4,87 @@
 # found in the LICENSE file.
 
 '''
-This script stitches the NetLog files in a specified directory.
+This script "stitches" the NetLog files from a ".inprogress" directory to
+create a single NetLog file.
+'''
+
+import glob
+import os
+import re
+import sys
+
+
+USAGE ='''Usage: stitch_net_log_files.py <INPROGRESS_DIR> [<OUTPUT_PATH>]
+
+Will copy all the files in <INPROGRESS_DIR> and write the their content into a
+NetLog file at path <OUTPUT_PATH>.
 
-The complete NetLog will be written to net-internals-log.json in the directory
-passed as argument to --path.
+If <OUTPUT_PATH> is not specified, it should end with ".inprogress", and the
+completed NetLog file will be written to the location with ".inprogress"
+stripped.
 '''
 
-import argparse, os
+
+def get_event_file_sort_key(path):
+  '''Returns a tuple (modification timestamp, file number) for a path of the
+  form event_file_%d.json'''
+
+  m = re.match('^event_file_(\d+).json$', path)
+  file_index = int(m.group(1))
+  return (os.path.getmtime(path), file_index)
+
+
+def get_ordered_event_files():
+  '''Returns a list of file paths to event files. The order of the files is
+  from oldest to newest. If modification times are the same, files will be
+  ordered based on the numeral in their file name.'''
+
+  paths = glob.glob("event_file_*.json")
+  paths = sorted(paths, key=get_event_file_sort_key)
+  sys.stdout.write("Identified %d event files:\n  %s\n" %
+                   (len(paths), "\n  ".join(paths)))
+  return paths
+
 
 def main():
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--path', action='store',
-      help="Specifies the complete filepath of the directory where the log "
-      "files are located.")
-  # TODO(dconnol): Automatically pull all event files matching the format
-  # event_file_<num>.json and remove the num_files argument.
-  parser.add_argument('--num_files', action='store',
-      help="Specifies the number of event files (not including the constants "
-      "file or the end_netlog file) that need need to be stitched together. "
-      "The number of event files passed to the script must not be greater "
-      "than the number of event files in the directory.")
-  args = parser.parse_args()
-
-  num_files = int(args.num_files)
-  filepath = args.path
-  if filepath[-1:] != "/":
-    filepath += "/"
-
-  os.chdir(filepath)
-
-  with open("net-internals-log.json", "w") as stitched_file:
+  if len(sys.argv) != 2 and len(sys.argv) != 3:
+    sys.stderr.write(USAGE)
+    sys.exit(1)
+
+  inprogress_dir = sys.argv[1]
+  output_path = None
+
+  # Pick an output path based on command line arguments.
+  if len(sys.argv) == 3:
+    output_path = sys.argv[2]
+  elif len(sys.argv) == 2:
+    m = re.match("^(.*)\.inprogress/?$", inprogress_dir)
+    if not m:
+      sys.stdout.write("Must specify OUTPUT_PATH\n")
+      sys.exit(1)
+    output_path = m.group(1)
+
+  output_path = os.path.abspath(output_path)
+
+  sys.stdout.write("Reading data from: %s\n" % inprogress_dir)
+  sys.stdout.write("Writing log file to: %s\n" % output_path)
+
+  os.chdir(inprogress_dir)
+
+  with open(output_path, "w") as stitched_file:
     try:
       file = open("constants.json")
       with file:
         for line in file:
           stitched_file.write(line)
     except IOError:
-      os.remove("net-internals-log.json")
-      print "File \"constants.json\" not found."
-      return
+      sys.stderr.write("Failed reading \"constants.json\".\n")
+      sys.exit(1)
 
     events_written = False;
-    for i in range(num_files):
+    for event_file_path in get_ordered_event_files():
       try:
-        file = open("event_file_%d.json" % i)
+        file = open(event_file_path)
         with file:
           if not events_written:
             line = file.readline();
@@ -59,9 +96,8 @@ def main():
               stitched_file.write(line)
               line = next_line
       except IOError:
-        os.remove("net-internals-log.json")
-        print "File \"event_file_%d.json\" not found." % i
-        return
+        sys.stderr.write("Failed reading \"%s\"\n" % event_file_path)
+        sys.exit(1)
     # Remove hanging comma from last event
     # TODO(dconnol): Check if the last line is a valid JSON object. If not,
     # do not write the line to file. This handles incomplete logs.
@@ -71,21 +107,21 @@ def main():
     elif line:
       raise ValueError('Last event is not properly formed')
 
+    if os.path.exists("end_netlog.json"):
       try:
         file = open("end_netlog.json")
         with file:
           for line in file:
             stitched_file.write(line)
       except IOError:
-        os.remove("net-internals-log.json")
-        print "File \"end_netlog\" not found."
-        return
-
-  # Delete old NetLog files
-  for i in range (num_files):
-    os.remove("event_file_%d.json" % i)
-  os.remove("constants.json")
-  os.remove("end_netlog.json")
+          sys.stderr.write("Failed reading \"end_netlog.json\".\n")
+          sys.exit(1)
+    else:
+      # end_netlog.json won't exist when using this tool to stitch logging
+      # sessions that didn't shutdown gracefully.
+      #
+      # Close the events array and then the log (no polled_data).
+      stitched_file.write("]}\n")
 
 
 if __name__ == "__main__":