Pass in custom options to Systrace agents

systrace/systrace/run_systrace.py

Index: systrace/systrace/run_systrace.py
diff --git a/systrace/systrace/run_systrace.py b/systrace/systrace/run_systrace.py
index 1d27c8eab9ecae329cb069dc699ea88a32441c2f..acab4642a5c77125e9aa62a5aee908112449f0a1 100755
--- a/systrace/systrace/run_systrace.py
+++ b/systrace/systrace/run_systrace.py
@@ -42,8 +42,14 @@ from devil import devil_env
 from devil.android.sdk import adb_wrapper
 from systrace import systrace_runner
 from systrace.tracing_agents import atrace_agent
+from systrace.tracing_agents import atrace_from_file_agent
+from systrace.tracing_agents import battor_trace_agent
 from systrace.tracing_agents import ftrace_agent
 
+ALL_MODULES = [atrace_agent, atrace_from_file_agent,
+               battor_trace_agent, ftrace_agent]
+
+
 def parse_options(argv):
   """Parses and checks the command-line options.
 
@@ -58,58 +64,17 @@ def parse_options(argv):
                     default=None, metavar='FILE')
   parser.add_option('-t', '--time', dest='trace_time', type='int',
                     help='trace for N seconds', metavar='N')
-  parser.add_option('-b', '--buf-size', dest='trace_buf_size', type='int',
-                    help='use a trace buffer size of N KB', metavar='N')
-  parser.add_option('-k', '--ktrace', dest='kfuncs', action='store',
-                    help='specify a comma-separated list of kernel functions '
-                    'to trace')
   parser.add_option('-l', '--list-categories', dest='list_categories',
                     default=False, action='store_true',
                     help='list the available categories and exit')
   parser.add_option('-j', '--json', dest='write_json',
                     default=False, action='store_true',
                     help='write a JSON file')
-  parser.add_option('-a', '--app', dest='app_name', default=None, type='string',
-                    action='store',
-                    help='enable application-level tracing for comma-separated '
-                    'list of app cmdlines')
-  parser.add_option('--no-fix-threads', dest='fix_threads', default=True,
-                    action='store_false',
-                    help='don\'t fix missing or truncated thread names')
-  parser.add_option('--no-fix-tgids', dest='fix_tgids', default=True,
-                    action='store_false',
-                    help='Do not run extra commands to restore missing thread '
-                    'to thread group id mappings.')
-  parser.add_option('--no-fix-circular', dest='fix_circular', default=True,
-                    action='store_false',
-                    help='don\'t fix truncated circular traces')
-  parser.add_option('--no-compress', dest='compress_trace_data',
-                    default=True, action='store_false',
-                    help='Tell the device not to send the trace data in '
-                    'compressed form.')
-  parser.add_option('--hubs', dest='hub_types', default='plugable_7port',
-                    help='List of hub types to check for for BattOr mapping. '
-                    'Used when updating mapping file.')
-  parser.add_option('--serial-map', dest='serial_map',
-                    default='serial_map.json',
-                    help='File containing pregenerated map of phone serial '
-                    'numbers to BattOr serial numbers.')
-  parser.add_option('--battor_path', dest='battor_path', default=None,
-                    type='string', help='specify a BattOr path to use')
-  parser.add_option('--update-map', dest='update_map', default=False,
-                    action='store_true',
-                    help='force update of phone-to-BattOr map')
   parser.add_option('--link-assets', dest='link_assets', default=False,
                     action='store_true',
                     help='(deprecated)')
-  parser.add_option('--boot', dest='boot', default=False, action='store_true',
-                    help='reboot the device with tracing during boot enabled. '
-                    'The report is created by hitting Ctrl+C after the device '
-                    'has booted up.')
-  parser.add_option('--battor', dest='battor', default=False,
-                    action='store_true', help='Use the BattOr tracing agent.')
   parser.add_option('--from-file', dest='from_file', action='store',
-                    help='read the trace from a file (compressed) rather than '
+                    help='read the trace from a file (compressed) rather than'
                     'running a live trace')
   parser.add_option('--asset-dir', dest='asset_dir', default='trace-viewer',
                     type='string', help='(deprecated)')
@@ -122,6 +87,33 @@ def parse_options(argv):
   parser.add_option('--collection-timeout', dest='collection_timeout',
                     type='int', help='timeout for data collection (seconds)')
 
+  atrace_ftrace_options = optparse.OptionGroup(parser,
+                                               'Atrace and Ftrace options')
+  atrace_ftrace_options.add_option('-b', '--buf-size', dest='trace_buf_size',
+                                   type='int', help='use a trace buffer size '
+                                   ' of N KB', metavar='N')
+  atrace_ftrace_options.add_option('--no-fix-threads', dest='fix_threads',
+                                   default=True, action='store_false',
+                                   help='don\'t fix missing or truncated '
+                                   'thread names')
+  atrace_ftrace_options.add_option('--no-fix-tgids', dest='fix_tgids',
+                                   default=True, action='store_false',
+                                   help='Do not run extra commands to restore'
+                                   ' missing thread to thread group id '
+                                   'mappings.')
+  atrace_ftrace_options.add_option('--no-fix-circular', dest='fix_circular',
+                                   default=True, action='store_false',
+                                   help='don\'t fix truncated circular traces')
+  parser.add_option_group(atrace_ftrace_options)
+
+  # Add the other agent parsing options to the parser. For Systrace on the
+  # command line, all agents are added. For Android, only the compatible agents
+  # will be added.
+  for module in ALL_MODULES:
+    option_group = module.add_options(parser)
+    if option_group:
+      parser.add_option_group(option_group)
+
   options, categories = parser.parse_args(argv[1:])
 
   if options.output_file is None:
@@ -164,6 +156,17 @@ def main():
   # Parse the command line options.
   options, categories = parse_options(sys.argv)
 
+  # Override --atrace-categories and --ftrace-categories flags if command-line
+  # categories are provided.
+  if categories:
+    if options.target == 'android':
+      options.atrace_categories = categories
+    elif options.target == 'linux':
+      options.ftrace_categories = categories
+    else:
+      raise RuntimeError('Categories are only valid for atrace/ftrace. Target '
+                         'platform must be either Android or Linux.')
+
   initialize_devil()
 
   if options.target == 'android' and not options.device_serial_number:
@@ -184,9 +187,8 @@ def main():
     return
 
   # Set up the systrace runner and start tracing.
-  script_dir = os.path.dirname(os.path.abspath(__file__))
   controller = systrace_runner.SystraceRunner(
-      script_dir, options, categories)
+      os.path.dirname(os.path.abspath(__file__)), options)
   controller.StartTracing()
 
   # Wait for the given number of seconds or until the user presses enter.