Add ability to capture SKPs from Chromium's pagesets

tools/skp/webpages_playback.py

 #!/usr/bin/env python
 # Copyright (c) 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.
 
 """Archives or replays webpages and creates SKPs in a Google Storage location.
 
 To archive webpages and store SKP files (archives should be rarely updated):
 
 cd skia
@@ skipping 97 matching lines @@
 CREDENTIALS_GS_PATH = '/playback/credentials/credentials.json'
 
 X11_DISPLAY = os.getenv('DISPLAY', ':0')
 
 GS_PREDEFINED_ACL = gs_utils.GSUtils.PredefinedACL.PRIVATE
 GS_FINE_GRAINED_ACL_LIST = [
   (gs_utils.GSUtils.IdType.GROUP_BY_DOMAIN, 'google.com',
    gs_utils.GSUtils.Permission.READ),
 ]
 
+# Path to Chromium's page sets.
+CHROMIUM_PAGE_SETS_PATH = os.path.join('tools', 'perf', 'page_sets')
+
+# Dictionary of supported Chromium page sets to their file prefixes.
+CHROMIUM_PAGE_SETS_TO_PREFIX = {
+    # 'key_mobile_sites_smooth.py': 'keymobi',
+    # TODO(rmistry): Uncomment the below after it is verified that the above
+    # works.
+    # 'top_25_smooth.py': 'top25desk',
+}
+
+
 def remove_prefix(s, prefix):
   if s.startswith(prefix):
     return s[len(prefix):]
   return s
 
+
 class SkPicturePlayback(object):
   """Class that archives or replays webpages and creates SKPs."""
 
   def __init__(self, parse_options):
     """Constructs a SkPicturePlayback BuildStep instance."""
     assert parse_options.browser_executable, 'Must specify --browser_executable'
     self._browser_executable = parse_options.browser_executable
 
+    self._chrome_page_sets_path = os.path.join(parse_options.chrome_src_path,
+                                               CHROMIUM_PAGE_SETS_PATH)
     self._all_page_sets_specified = parse_options.page_sets == 'all'
     self._page_sets = self._ParsePageSets(parse_options.page_sets)
 
     self._record = parse_options.record
     self._skia_tools = parse_options.skia_tools
     self._non_interactive = parse_options.non_interactive
     self._upload = parse_options.upload
     data_store_location = parse_options.data_store
     if data_store_location.startswith(gs_utils.GS_PREFIX):
       self.gs = GoogleStorageDataStore(data_store_location)
@@ skipping 15 matching lines @@
     if not page_sets:
       raise ValueError('Must specify at least one page_set!')
     elif self._all_page_sets_specified:
       # Get everything from the page_sets directory.
       page_sets_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)),
                                    'page_sets')
       ps = [os.path.join(page_sets_dir, page_set)
             for page_set in os.listdir(page_sets_dir)
             if not os.path.isdir(os.path.join(page_sets_dir, page_set)) and
                page_set.endswith('.py')]
+      chromium_ps = [
+          os.path.join(self._chrome_page_sets_path, cr_page_set)
+          for cr_page_set in CHROMIUM_PAGE_SETS_TO_PREFIX]
+      ps.extend(chromium_ps)
     elif '*' in page_sets:
       # Explode and return the glob.
       ps = glob.glob(page_sets)
     else:
       ps = page_sets.split(',')
     ps.sort()
     return ps
 
+  def _IsChromiumPageSet(self, page_set):
+    """Returns true if the specified page set is a Chromium page set."""
+    return page_set.startswith(self._chrome_page_sets_path)
+
   def Run(self):
     """Run the SkPicturePlayback BuildStep."""
 
     # Download the credentials file if it was not previously downloaded.
     if not os.path.isfile(CREDENTIALS_FILE_PATH):
       # Download the credentials.json file from Google Storage.
       self.gs.download_file(CREDENTIALS_GS_PATH, CREDENTIALS_FILE_PATH)
 
     if not os.path.isfile(CREDENTIALS_FILE_PATH):
       print """\n\nCould not locate credentials file in the storage.
@@ skipping 26 matching lines @@
     start_time = time.time()
 
     # Loop through all page_sets.
     for page_set in self._page_sets:
 
       page_set_basename = os.path.basename(page_set).split('.')[0]
       page_set_json_name = page_set_basename + '.json'
       wpr_data_file = page_set.split(os.path.sep)[-1].split('.')[0] + '_000.wpr'
       page_set_dir = os.path.dirname(page_set)
 
-      if self._record:
+      if self._IsChromiumPageSet(page_set):
+        print 'Using Chromium\'s captured archives for Chromium\'s page sets.'
+      elif self._record:
         # Create an archive of the specified webpages if '--record=True' is
         # specified.
         record_wpr_cmd = (
           'PYTHONPATH=%s:$PYTHONPATH' % page_set_dir,
           'DISPLAY=%s' % X11_DISPLAY,
           os.path.join(self._telemetry_binaries_dir, 'record_wpr'),
           '--extra-browser-args=--disable-setuid-sandbox',
           '--browser=exact',
           '--browser-executable=%s' % self._browser_executable,
           '%s_page_set' % page_set_basename,
@@ skipping 109 matching lines @@
       print '\n\n=======New SKPs have been uploaded to %s =======\n\n' % (
           posixpath.join(self.gs.target_name(), dest_dir_name,
                          SKPICTURES_DIR_NAME))
     else:
       print '\n\n=======Not Uploading to %s=======\n\n' % self.gs.target_type()
       print 'Generated resources are available in %s\n\n' % (
           LOCAL_PLAYBACK_ROOT_DIR)
 
     return 0
 
+  def _GetSkiaSkpFileName(self, page_set):
+    """Returns the SKP file name for Skia page sets."""
+    # /path/to/skia_yahooanswers_desktop.py -> skia_yahooanswers_desktop.py
+    ps_filename = os.path.basename(page_set)
+    # skia_yahooanswers_desktop.py -> skia_yahooanswers_desktop
+    ps_basename, _ = os.path.splitext(ps_filename)
+    # skia_yahooanswers_desktop -> skia, yahooanswers, desktop
+    _, page_name, device = ps_basename.split('_')
+    basename = '%s_%s' % (DEVICE_TO_PLATFORM_PREFIX[device], page_name)
+    return basename[:MAX_SKP_BASE_NAME_LEN] + '.skp'
+
+  def _GetChromiumSkpFileName(self, page_set, site):
+    """Returns the SKP file name for Chromium page sets."""
+    # /path/to/http___mobile_news_sandbox_pt0 -> http___mobile_news_sandbox_pt0
+    _, webpage = os.path.split(site)
+    # http___mobile_news_sandbox_pt0 -> mobile_news_sandbox_pt0
+    webpage = webpage.lstrip('http___').lstrip('https___')
+    # /path/to/skia_yahooanswers_desktop.py -> skia_yahooanswers_desktop.py
+    ps_filename = os.path.basename(page_set)
+    # http___mobile_news_sandbox -> pagesetprefix_http___mobile_news_sandbox
+    basename = '%s_%s' % (CHROMIUM_PAGE_SETS_TO_PREFIX[ps_filename], webpage)
+    return basename[:MAX_SKP_BASE_NAME_LEN] + '.skp'
+
   def _RenameSkpFiles(self, page_set):
     """Rename generated SKP files into more descriptive names.
 
     Look into the subdirectory of TMP_SKP_DIR and find the most interesting
     .skp in there to be this page_set's representative .skp.
     """
-    # Here's where we're assuming there's one page per pageset.
-    # If there were more than one, we'd overwrite filename below.
+    subdirs = glob.glob(os.path.join(TMP_SKP_DIR, '*'))
+    for site in subdirs:
+      if self._IsChromiumPageSet(page_set):
+        filename = self._GetChromiumSkpFileName(page_set, site)
+      else:
+        filename = self._GetSkiaSkpFileName(page_set)
 
-    # /path/to/skia_yahooanswers_desktop.json -> skia_yahooanswers_desktop.json
-    _, ps_filename = os.path.split(page_set)
-    # skia_yahooanswers_desktop.json -> skia_yahooanswers_desktop
-    ps_basename, _ = os.path.splitext(ps_filename)
-    # skia_yahooanswers_desktop -> skia, yahooanswers, desktop
-    _, page_name, device = ps_basename.split('_')
-
-    basename = '%s_%s' % (DEVICE_TO_PLATFORM_PREFIX[device], page_name)
-    filename = basename[:MAX_SKP_BASE_NAME_LEN] + '.skp'
-
-    subdirs = glob.glob(os.path.join(TMP_SKP_DIR, '*'))
-    assert len(subdirs) == 1
-    for site in subdirs:
       # We choose the largest .skp as the most likely to be interesting.
       largest_skp = max(glob.glob(os.path.join(site, '*.skp')),
                         key=lambda path: os.stat(path).st_size)
       dest = os.path.join(self._local_skp_dir, filename)
       print 'Moving', largest_skp, 'to', dest
       shutil.move(largest_skp, dest)
       self._skp_files.append(filename)
       shutil.rmtree(site)
 
   def _CreateLocalStorageDirs(self):
@@ skipping 34 matching lines @@
   def does_storage_object_exist(self, *args):
     raise NotImplementedError()
   def download_file(self, *args):
     raise NotImplementedError()
   def upload_dir_contents(self, source_dir, **kwargs):
     raise NotImplementedError()
 
 class GoogleStorageDataStore(DataStore):
   def __init__(self, data_store_url):
     self._data_store_url = data_store_url
-    self._bucket = remove_prefix(self._data_store_url.lstrip(), 
+    self._bucket = remove_prefix(self._data_store_url.lstrip(),
                                  gs_utils.GS_PREFIX)
     self.gs = gs_utils.GSUtils()
   def target_name(self):
     return self._data_store_url
   def target_type(self):
     return 'Google Storage'
   def does_storage_object_exist(self, *args):
     return self.gs.does_storage_object_exist(self._bucket, *args)
   def download_file(self, *args):
     self.gs.download_file(self._bucket, *args)
@@ skipping 74 matching lines @@
       default=None)
   option_parser.add_option(
       '', '--non-interactive', action='store_true',
       help='Runs the script without any prompts. If this flag is specified and '
            '--skia_tools is specified then the debugger is not run.',
       default=False)
   options, unused_args = option_parser.parse_args()
 
   playback = SkPicturePlayback(options)
   sys.exit(playback.Run())