Add usage message with available page_sets and benchmarks.

tools/telemetry/telemetry/page/record_wpr.py

 # Copyright 2012 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.
 
 import logging
 import sys
 
 from telemetry import benchmark
 from telemetry.core import browser_options
 from telemetry.core import discover
@@ skipping 85 matching lines @@
     if not issubclass(cls, profile_creator.ProfileCreator):
       action_name = cls().action_name_to_run
       if action_name:
         action_names.add(action_name)
   return action_names
 
 
 def _MaybeGetInstanceOfClass(target, base_dir, cls):
   if isinstance(target, cls):
     return target
-  classes = discover.DiscoverClasses(base_dir, base_dir, cls,
-                                     index_by_class_name=True)
-  return classes[target]() if target in classes else None
+  matching_classes = discover.MatchName([cls], target, base_dir=base_dir)
+  class_instance = None
+  if len(matching_classes) == 1:
+    try:
+      class_instance = matching_classes[0]()
+    except TypeError:
+      logging.warn('Unable to initialize "%s" class which requires ' % target +
+                   'custom arguments.')
+      raise
+  return class_instance
 
 
 class WprRecorder(object):
 
-  def __init__(self, base_dir, target, args=None):
+  def __init__(self, base_dir, target, cmd, args=None):
     action_names_to_run = FindAllActionNames(base_dir)
     self._record_page_test = RecorderPageTest(action_names_to_run)
     self._options = self._CreateOptions()
+    self._page_set = None
+    self._benchmark = None
 
-    self._benchmark = _MaybeGetInstanceOfClass(target, base_dir,
-                                               benchmark.Benchmark)
-    if self._benchmark is not None:
-      self._record_page_test.page_test = self._benchmark.test()
-    self._parser = self._options.CreateParser(usage='%prog <PageSet|Benchmark>')
+    if cmd == 'run' or cmd == 'run_pageset':
+      self._page_set = _MaybeGetInstanceOfClass(
+          target, base_dir, page_set.PageSet)
+    if self._page_set is None and (cmd == 'run' or cmd =='run_benchmark'):
+      self._benchmark = _MaybeGetInstanceOfClass(
+          target, base_dir, benchmark.Benchmark)
+    if self._page_set is None and self._benchmark is None:
+      print_usage()
+      sys.exit(1)
+    recording_type = 'page set' if self._benchmark is None else 'benchmark'
+    print 'Recording %s: %s' % (recording_type, target)
+
+    self._parser = self._options.CreateParser()
     self._AddCommandLineArgs()
     self._ParseArgs(args)
     self._ProcessCommandLineArgs()
-    self._page_set = self._GetPageSet(base_dir, target)
+
+    if self._page_set is None:
+      self._record_page_test.page_test = self._benchmark.test()
+      self._page_set = self._benchmark.CreatePageSet(self._options)
 
   @property
   def options(self):
     return self._options
 
   def _CreateOptions(self):
     options = browser_options.BrowserFinderOptions()
     options.browser_options.wpr_mode = wpr_modes.WPR_RECORD
     options.browser_options.no_proxy_server = True
     return options
@@ skipping 18 matching lines @@
 
   def _ParseArgs(self, args=None):
     args_to_parse = sys.argv[1:] if args is None else args
     self._parser.parse_args(args_to_parse)
 
   def _ProcessCommandLineArgs(self):
     page_runner.ProcessCommandLineArgs(self._parser, self._options)
     if self._benchmark is not None:
       self._benchmark.ProcessCommandLineArgs(self._parser, self._options)
 
-  def _GetPageSet(self, base_dir, target):
-    if self._benchmark is not None:
-      return self._benchmark.CreatePageSet(self._options)
-    ps = _MaybeGetInstanceOfClass(target, base_dir, page_set.PageSet)
-    if ps is None:
-      self._parser.print_usage()
-      sys.exit(1)
-    return ps
-
   def Record(self, results):
     self._page_set.wpr_archive_info.AddNewTemporaryRecording()
     self._record_page_test.CustomizeBrowserOptions(self._options)
     page_runner.Run(self._record_page_test, self._page_set,
         test_expectations.TestExpectations(), self._options, results)
 
   def HandleResults(self, results):
     if results.failures or results.skipped_values:
       logging.warning('Some pages failed and/or were skipped. The recording '
                       'has not been updated for these pages.')
     results.PrintSummary()
     self._page_set.wpr_archive_info.AddRecordedPages(
         results.pages_that_succeeded)
 
 
 def Main(base_dir):
   quick_args = [a for a in sys.argv[1:] if not a.startswith('-')]
-  if len(quick_args) != 1:
-    print >> sys.stderr, 'Usage: record_wpr <PageSet|Benchmark>\n'
-    sys.exit(1)
-  target = quick_args.pop()
-  wpr_recorder = WprRecorder(base_dir, target)
+  if len(quick_args) == 1:
+    arg = quick_args.pop()
+    if (arg == 'list_pagesets'):
+      return discover.PrintAvailableClasses([page_set.PageSet])
+    if (arg == 'list_benchmarks'):
+      return discover.PrintAvailableClasses([benchmark.Benchmark])
+
+    # Otherwise, assume no command is provided and default to 'run'
+    return run(base_dir, arg, 'run')
+  if len(quick_args) == 2:
+    cmd, target = quick_args
+    if (cmd == 'run' or cmd == 'run_benchmark' or cmd == 'run_pageset'):
+      return run(base_dir, target, cmd)
+  print_usage()
+  return 1
+
+
+def run(base_dir, target, cmd):
+  wpr_recorder = WprRecorder(base_dir, target, cmd)
   results = wpr_recorder.CreateResults()
   wpr_recorder.Record(results)
   wpr_recorder.HandleResults(results)
   return min(255, len(results.failures))
+
+
+def print_usage():
+  sys.stderr.write(
+      'Usage: record_wpr [command] <PageSet|Benchmark>\n'
+      'Available commands are:\n'
+      '  run <PageSet|Benchmark>       Run a Page Set or Benchmark (default).\n'
+      '                                In case of name conflicts, Page Sets\n'
+      '                                will take precedence over Benchmarks.\n'
+      '  run_benchmark <Benchmark>     Run a Benchmark\n'
+      '  run_pageset <PageSet>         Run a Page Set\n'
+      '  list_pagesets                 Lists the available Page Sets\n'
+      '  list_benchmarks               Lists the available Benchmarks\n\n')