sandwich: Makes pull_sandwich_metrics.py a sub-command.

tools/android/loading/pull_sandwich_metrics.py

-#! /usr/bin/env python
 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Pull a sandwich run's output directory's metrics from traces into a CSV.
 
 python pull_sandwich_metrics.py -h
 """
 
-import argparse
-import csv
 import json
 import logging
 import os
 import sys
 
 import loading_trace as loading_trace_module
 import tracing
 
 
 CATEGORIES = ['blink.user_timing', 'disabled-by-default-memory-infra']
 
-_CSV_FIELD_NAMES = [
+CSV_FIELD_NAMES = [
     'id',
     'url',
     'total_load',
     'onload',
     'browser_malloc_avg',
     'browser_malloc_max']
 
 _TRACKED_EVENT_NAMES = set(['requestStart', 'loadEventStart', 'loadEventEnd'])
 
 
@@ skipping 72 matching lines @@
   return tracked_events
 
 
 def _PullMetricsFromLoadingTrace(loading_trace):
   """Pulls all the metrics from a given trace.
 
   Args:
     loading_trace: loading_trace_module.LoadingTrace.
 
   Returns:
-    Dictionary with all _CSV_FIELD_NAMES's field set (except the 'id').
+    Dictionary with all CSV_FIELD_NAMES's field set (except the 'id').
   """
   browser_dump_events = _GetBrowserDumpEvents(loading_trace.tracing_track)
   web_page_tracked_events = _GetWebPageTrackedEvents(
       loading_trace.tracing_track)
 
   browser_malloc_sum = 0
   browser_malloc_max = 0
   for dump_event in browser_dump_events:
     attr = dump_event.args['dumps']['allocators']['malloc']['attrs']['size']
     assert attr['units'] == 'bytes'
     size = int(attr['value'], 16)
     browser_malloc_sum += size
     browser_malloc_max = max(browser_malloc_max, size)
 
   return {
     'total_load': (web_page_tracked_events['loadEventEnd'].start_msec -
                    web_page_tracked_events['requestStart'].start_msec),
     'onload': (web_page_tracked_events['loadEventEnd'].start_msec -
                web_page_tracked_events['loadEventStart'].start_msec),
     'browser_malloc_avg': browser_malloc_sum / float(len(browser_dump_events)),
     'browser_malloc_max': browser_malloc_max
   }
 
 
-def _PullMetricsFromOutputDirectory(output_directory_path):
+def PullMetricsFromOutputDirectory(output_directory_path):
   """Pulls all the metrics from all the traces of a sandwich run directory.
 
   Args:
     output_directory_path: The sandwich run's output directory to pull the
         metrics from.
 
   Returns:
-    List of dictionaries with all _CSV_FIELD_NAMES's field set.
+    List of dictionaries with all CSV_FIELD_NAMES's field set.
   """
   assert os.path.isdir(output_directory_path)
   run_infos = None
   with open(os.path.join(output_directory_path, 'run_infos.json')) as f:
     run_infos = json.load(f)
   assert run_infos
   metrics = []
   for node_name in os.listdir(output_directory_path):
     if not os.path.isdir(os.path.join(output_directory_path, node_name)):
       continue
     try:
       page_id = int(node_name)
     except ValueError:
       continue
     trace_path = os.path.join(output_directory_path, node_name, 'trace.json')
     if not os.path.isfile(trace_path):
       continue
     logging.info('processing \'%s\'' % trace_path)
     loading_trace = loading_trace_module.LoadingTrace.FromJsonFile(trace_path)
     trace_metrics = _PullMetricsFromLoadingTrace(loading_trace)
     trace_metrics['id'] = page_id
     trace_metrics['url'] = run_infos['urls'][page_id]
     metrics.append(trace_metrics)
   assert len(metrics) > 0, ('Looks like \'{}\' was not a sandwich ' +
                             'run directory.').format(output_directory_path)
   return metrics
-
-
-def main():
-  logging.basicConfig(level=logging.INFO)
-
-  parser = argparse.ArgumentParser()
-  parser.add_argument('output', type=str,
-                      help='Output directory of run_sandwich.py command.')
-  args = parser.parse_args()
-
-  trace_metrics_list = _PullMetricsFromOutputDirectory(args.output)
-  trace_metrics_list.sort(key=lambda e: e['id'])
-  cs_file_path = os.path.join(args.output, 'trace_analysis.csv')
-  with open(cs_file_path, 'w') as csv_file:
-    writer = csv.DictWriter(csv_file, fieldnames=_CSV_FIELD_NAMES)
-    writer.writeheader()
-    for trace_metrics in trace_metrics_list:
-      writer.writerow(trace_metrics)
-  return 0
-
-
-if __name__ == '__main__':
-  sys.exit(main())