webpages_playback.py: Support storing to local directories

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
-python tools/skp/webpages_playback.py --dest_gsbase=gs://rmistry --record \
+python tools/skp/webpages_playback.py --data_store=gs://rmistry --record \
 --page_sets=all --skia_tools=/home/default/trunk/out/Debug/ \
 --browser_executable=/tmp/chromium/out/Release/chrome
 
+The above command uses Google Storage bucket 'rmistry' to download needed files.
 
 To replay archived webpages and re-generate SKP files (should be run whenever
 SkPicture.PICTURE_VERSION changes):
 
 cd skia
-python tools/skp/webpages_playback.py --dest_gsbase=gs://rmistry \
+python tools/skp/webpages_playback.py --data_store=gs://rmistry \
 --page_sets=all --skia_tools=/home/default/trunk/out/Debug/ \
 --browser_executable=/tmp/chromium/out/Release/chrome
 
 
 Specify the --page_sets flag (default value is 'all') to pick a list of which
 webpages should be archived and/or replayed. Eg:
 
 --page_sets=tools/skp/page_sets/skia_yahooanswers_desktop.py,\
 tools/skp/page_sets/skia_googlecalendar_nexus10.py
 
 The --browser_executable flag should point to the browser binary you want to use
 to capture archives and/or capture SKP files. Majority of the time it should be
 a newly built chrome binary.
 
-The --upload_to_gs flag controls whether generated artifacts will be uploaded
-to Google Storage (default value is False if not specified).
+The --data_store flag controls where the needed artifacts, such as
+credential files, are downloaded from. It also controls where the
+generated artifacts, such as recorded webpages and resulting skp renderings,
+are uploaded to. URLs with scheme 'gs://' use Google Storage. Otherwise
+use local filesystem.
+
+The --upload=True flag means generated artifacts will be
+uploaded or copied to the location specified by --data_store. (default value is
+False if not specified).
 
 The --non-interactive flag controls whether the script will prompt the user
 (default value is False if not specified).
 
 The --skia_tools flag if specified will allow this script to run
 debugger, render_pictures, and render_pdfs on the captured
 SKP(s). The tools are run after all SKPs are succesfully captured to make sure
 they can be added to the buildbots with no breakages.
 """
 
@@ skipping 69 matching lines @@
   """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._all_page_sets_specified = parse_options.page_sets == 'all'
     self._page_sets = self._ParsePageSets(parse_options.page_sets)
 
-    self._dest_gsbase = parse_options.dest_gsbase
     self._record = parse_options.record
     self._skia_tools = parse_options.skia_tools
     self._non_interactive = parse_options.non_interactive
-    self._upload_to_gs = parse_options.upload_to_gs
+    self._upload = parse_options.upload
+    data_store_location = parse_options.data_store
+    if data_store_location.startswith(gs_utils.GS_PREFIX):
+      self.gs = GSWithBucket(data_store_location)
+    else:
+      self.gs = GSFileEmulation(data_store_location)
     self._alternate_upload_dir = parse_options.alternate_upload_dir
-    self._skip_all_gs_access = parse_options.skip_all_gs_access
     self._telemetry_binaries_dir = os.path.join(parse_options.chrome_src_path,
                                                 'tools', 'perf')
 
     self._local_skp_dir = os.path.join(
         parse_options.output_dir, ROOT_PLAYBACK_DIR_NAME, SKPICTURES_DIR_NAME)
     self._local_record_webpages_archive_dir = os.path.join(
         parse_options.output_dir, ROOT_PLAYBACK_DIR_NAME, 'webpages_archive')
 
     # List of SKP files generated by this script.
     self._skp_files = []
@@ skipping 14 matching lines @@
       ps = glob.glob(page_sets)
     else:
       ps = page_sets.split(',')
     ps.sort()
     return ps
 
   def Run(self):
     """Run the SkPicturePlayback BuildStep."""
 
     # Download the credentials file if it was not previously downloaded.
-    if self._skip_all_gs_access:
-      print """\n\nPlease create a %s file that contains:
+    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.
+      Please create a %s file that contains:
       {
         "google": {
           "username": "google_testing_account_username",
           "password": "google_testing_account_password"
         },
         "facebook": {
           "username": "facebook_testing_account_username",
           "password": "facebook_testing_account_password"
         }
       }\n\n""" % CREDENTIALS_FILE_PATH
       raw_input("Please press a key when you are ready to proceed...")
-    elif not os.path.isfile(CREDENTIALS_FILE_PATH):
-      # Download the credentials.json file from Google Storage.
-      gs_bucket = remove_prefix(self._dest_gsbase.lstrip(), gs_utils.GS_PREFIX)
-      gs_utils.GSUtils().download_file(gs_bucket, CREDENTIALS_GS_PATH,
-                                       CREDENTIALS_FILE_PATH)
 
     # Delete any left over data files in the data directory.
     for archive_file in glob.glob(
         os.path.join(LOCAL_REPLAY_WEBPAGES_ARCHIVE_DIR, 'skia_*')):
       os.remove(archive_file)
 
     # Delete the local root directory if it already exists.
     if os.path.exists(LOCAL_PLAYBACK_ROOT_DIR):
       shutil.rmtree(LOCAL_PLAYBACK_ROOT_DIR)
 
@@ skipping 31 matching lines @@
             break
           except Exception:
             # There was a failure continue with the loop.
             traceback.print_exc()
         else:
           # If we get here then record_wpr did not succeed and thus did not
           # break out of the loop.
           raise Exception('record_wpr failed for page_set: %s' % page_set)
 
       else:
-        if not self._skip_all_gs_access:
-          # Get the webpages archive so that it can be replayed.
-          self._DownloadWebpagesArchive(wpr_data_file, page_set_json_name)
+        # Get the webpages archive so that it can be replayed.
+        self._DownloadWebpagesArchive(wpr_data_file, page_set_json_name)
 
       run_benchmark_cmd = (
           'PYTHONPATH=%s:$PYTHONPATH' % page_set_dir,
           'DISPLAY=%s' % X11_DISPLAY,
           'timeout', '300',
           os.path.join(self._telemetry_binaries_dir, 'run_benchmark'),
           '--extra-browser-args=--disable-setuid-sandbox',
           '--browser=exact',
           '--browser-executable=%s' % self._browser_executable,
           SKP_BENCHMARK,
@@ skipping 58 matching lines @@
         if code != 0:
           raise Exception('%s failed!' % ' '.join(tools_cmd))
 
       if not self._non_interactive:
         print '\n\n=======Running debugger======='
         os.system('%s %s' % (os.path.join(self._skia_tools, 'debugger'),
                              self._local_skp_dir))
 
     print '\n\n'
 
-    if not self._skip_all_gs_access and self._upload_to_gs:
-      print '\n\n=======Uploading to Google Storage=======\n\n'
+    if self._upload:
+      print '\n\n=======Uploading to %s=======\n\n' % self.gs.target_type()
       # Copy the directory structure in the root directory into Google Storage.
       dest_dir_name = ROOT_PLAYBACK_DIR_NAME
       if self._alternate_upload_dir:
         dest_dir_name = self._alternate_upload_dir
 
-      gs_bucket = remove_prefix(self._dest_gsbase.lstrip(), gs_utils.GS_PREFIX)
-      gs_utils.GSUtils().upload_dir_contents(
-          LOCAL_PLAYBACK_ROOT_DIR, gs_bucket, dest_dir_name,
+      self.gs.upload_dir_contents(
+          LOCAL_PLAYBACK_ROOT_DIR, dest_dir_name,
           upload_if=gs_utils.GSUtils.UploadIf.IF_MODIFIED,
           predefined_acl=GS_PREDEFINED_ACL,
           fine_grained_acl_list=GS_FINE_GRAINED_ACL_LIST)
 
       print '\n\n=======New SKPs have been uploaded to %s =======\n\n' % (
-          posixpath.join(self._dest_gsbase, dest_dir_name, SKPICTURES_DIR_NAME))
+          posixpath.join(self.gs.target_name(), dest_dir_name,
+                         SKPICTURES_DIR_NAME))
     else:
-      print '\n\n=======Not Uploading to Google Storage=======\n\n'
+      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 _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.
@@ skipping 31 matching lines @@
         shutil.rmtree(d)
       os.makedirs(d)
 
   def _DownloadWebpagesArchive(self, wpr_data_file, page_set_json_name):
     """Downloads the webpages archive and its required page set from GS."""
     wpr_source = posixpath.join(ROOT_PLAYBACK_DIR_NAME, 'webpages_archive',
                                 wpr_data_file)
     page_set_source = posixpath.join(ROOT_PLAYBACK_DIR_NAME,
                                      'webpages_archive',
                                      page_set_json_name)
-    gs = gs_utils.GSUtils()
-    gs_bucket = remove_prefix(self._dest_gsbase.lstrip(), gs_utils.GS_PREFIX)
-    if (gs.does_storage_object_exist(gs_bucket, wpr_source) and
-        gs.does_storage_object_exist(gs_bucket, page_set_source)):
-      gs.download_file(gs_bucket, wpr_source,
+    gs = self.gs
+    if (gs.does_storage_object_exist(wpr_source) and
+        gs.does_storage_object_exist(page_set_source)):
+      gs.download_file(wpr_source,
                        os.path.join(LOCAL_REPLAY_WEBPAGES_ARCHIVE_DIR,
                                     wpr_data_file))
-      gs.download_file(gs_bucket, page_set_source,
+      gs.download_file(page_set_source,
                        os.path.join(LOCAL_REPLAY_WEBPAGES_ARCHIVE_DIR,
                                     page_set_json_name))
     else:
-      raise Exception('%s and %s do not exist in Google Storage!' % (
-          wpr_source, page_set_source))
+      raise Exception('%s and %s do not exist in %s!' % (gs.target_type(),
+        wpr_source, page_set_source))
 
+class GSWithBucket:

[rmistry, 2014/12/29 12:28:59]

The following is just a suggestion let me know if you think it might be
overkill:

Could you create an abstract class for these two classes to inherit from?
something generic like DataStore and the subclasses can be something like
LocalFileSystem and GoogleStorage (or better names).
The abstract class can probably implement target_name and target_type and throw
exceptions for the ones it expects subclasses to implement.

+  def __init__(self, data_store_url):
+    self._data_store_url = data_store_url
+    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)
+  def upload_dir_contents(self, source_dir, **kwargs):
+    self.gs.upload_dir_contents(source_dir, self._bucket, **kwargs)
+
+class GSFileEmulation:
+  def __init__(self, data_store_location):
+    self._base_dir = data_store_location
+  def target_name(self):
+    return self._base_dir;
+  def target_type(self):
+    return self._base_dir
+  def does_storage_object_exist(self, name):
+    return os.path.isfile(os.path.join(self._base_dir, name))
+  def download_file(self, name, local_path):
+    shutil.copyfile(os.path.join(self._base_dir, name), local_path)
+  def upload_dir_contents(self, source_dir, dest_dir, **kwargs):
+    def copytree(source_dir, dest_dir):
+      if not os.path.exists(dest_dir):
+        os.makedirs(dest_dir)
+      for item in os.listdir(source_dir):
+        source = os.path.join(source_dir, item)
+        dest = os.path.join(dest_dir, item)
+        if os.path.isdir(source):
+          copytree(source, dest)
+        else:
+          shutil.copy2(source, dest)
+    copytree(source_dir, os.path.join(self._base_dir, dest_dir))
 
 if '__main__' == __name__:
   option_parser = optparse.OptionParser()
   option_parser.add_option(
       '', '--page_sets',
       help='Specifies the page sets to use to archive. Supports globs.',
       default='all')
   option_parser.add_option(
-      '', '--skip_all_gs_access', action='store_true',
-      help='All Google Storage interactions will be skipped if this flag is '
-           'specified. This is useful for cases where the user does not have '
-           'the required .boto file but would like to generate webpage '
-           'archives and SKPs from the Skia page sets.',
-      default=False)
-  option_parser.add_option(
       '', '--record', action='store_true',
       help='Specifies whether a new website archive should be created.',
       default=False)
   option_parser.add_option(
-      '', '--dest_gsbase',
-      help='gs:// bucket_name, the bucket to upload the file to.',
-      default='gs://chromium-skia-gm')
-  option_parser.add_option(
       '', '--skia_tools',
       help=('Path to compiled Skia executable tools. '
             'render_pictures/render_pdfs is run on the set '
             'after all SKPs are captured. If the script is run without '
             '--non-interactive then the debugger is also run at the end. Debug '
             'builds are recommended because they seem to catch more failures '
             'than Release builds.'),
       default=None)
   option_parser.add_option(
-      '', '--upload_to_gs', action='store_true',
-      help='Does not upload to Google Storage if this is False.',
+      '', '--upload', action='store_true',
+      help=('Uploads to Google Storage or copies to local filesystem storage '
+            ' if this is True.'),
       default=False)
   option_parser.add_option(
+      '', '--data_store',
+    help=('The location of the file storage to use to download and upload '
+          'files. Can be \'gs://<bucket>\' for Google Storage, or '
+          'a directory for local filesystem storage'),
+      default='gs://chromium-skia-gm')
+  option_parser.add_option(
       '', '--alternate_upload_dir',
-      help='Uploads to a different directory in Google Storage if this flag is '
-           'specified',
+      help= ('Uploads to a different directory in Google Storage or local '
+             'storage if this flag is specified'),
       default=None)
   option_parser.add_option(
       '', '--output_dir',
-      help='Directory where SKPs and webpage archives will be outputted to.',
+      help=('Temporary directory where SKPs and webpage archives will be '
+            'outputted to.'),
       default=tempfile.gettempdir())
   option_parser.add_option(
       '', '--browser_executable',
       help='The exact browser executable to run.',
       default=None)
   option_parser.add_option(
       '', '--chrome_src_path',
       help='Path to the chromium src directory.',
       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())