Refactor downloading archives in archive_info.py

tools/telemetry/telemetry/wpr/archive_info.py

 # Copyright 2013 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 json
 import logging
 import os
 import re
 import shutil
 import tempfile
 
 from telemetry import page as page_module
 from telemetry.util import cloud_storage
 
 
 def AssertValidCloudStorageBucket(bucket):
   is_valid = bucket in (None,
                         cloud_storage.PUBLIC_BUCKET,
                         cloud_storage.PARTNER_BUCKET,
                         cloud_storage.INTERNAL_BUCKET)
   if not is_valid:
     raise ValueError("Cloud storage privacy bucket %s is invalid" % bucket)
 
 
+class ArchiveError(Exception):
+  pass
+
+
 class WprArchiveInfo(object):
-  def __init__(self, file_path, data, bucket, ignore_archive=False):
+  def __init__(self, file_path, data, bucket):
     AssertValidCloudStorageBucket(bucket)
     self._file_path = file_path
     self._base_dir = os.path.dirname(file_path)
+    self._data = data
     self._bucket = bucket
 
     # Ensure directory exists.
     if not os.path.exists(self._base_dir):
       os.makedirs(self._base_dir)
 
-    # TODO(aiolos): We should take this out of init if we reduce the number of
-    # supported code paths/configs using archive_info when switching over to
-    # user_stories.
-    # Download all .wpr files.
-    if not ignore_archive:
-      if not self._bucket:
-        logging.warning('User story set in %s has no bucket specified, and '
-                        'cannot be downloaded from cloud_storage.', file_path)
-      else:
-        for archive_path in data['archives']:
-          archive_path = self._WprFileNameToPath(archive_path)
-          try:
-            cloud_storage.GetIfChanged(archive_path, bucket)
-          except (cloud_storage.CredentialsError,
-                  cloud_storage.PermissionError):
-            if os.path.exists(archive_path):
-              # If the archive exists, assume the user recorded their own and
-              # simply warn.
-              logging.warning('Need credentials to update WPR archive: %s',
-                              archive_path)
-            else:
-              logging.error("You either aren't authenticated or don't have "
-                            "permission to use the archives for this page set."
-                            "\nYou may need to run gsutil config")
-              raise
-
     # Map from the relative path (as it appears in the metadata file) of the
     # .wpr file to a list of user story names it supports.
     self._wpr_file_to_user_story_names = data['archives']
 
     # Map from the user_story name to a relative path (as it appears
     # in the metadata file) of the .wpr file.
     self._user_story_name_to_wpr_file = dict()
     # Find out the wpr file names for each user_story.
     for wpr_file in data['archives']:
       user_story_names = data['archives'][wpr_file]
       for user_story_name in user_story_names:
         self._user_story_name_to_wpr_file[user_story_name] = wpr_file
     self.temp_target_wpr_file_path = None
 
   @classmethod
-  def FromFile(cls, file_path, bucket, ignore_archive=False):
+  def FromFile(cls, file_path, bucket):
     if os.path.exists(file_path):
       with open(file_path, 'r') as f:
         data = json.load(f)
-        return cls(file_path, data, bucket, ignore_archive=ignore_archive)
-    return cls(file_path, {'archives': {}}, bucket,
-               ignore_archive=ignore_archive)
+        return cls(file_path, data, bucket)
+    return cls(file_path, {'archives': {}}, bucket)
+
+  def DownloadArchivesIfNeeded(self):
+    """Downloads archives iff the Archive has a bucket parameter and the user
+    has permission to access the bucket.
+
+    Raises cloud storage Permissions or Credentials error when there is no
+    local copy of the archive and the user doesn't have permission to access
+    the archive's bucket.
+
+    Warns when a bucket is not specified or when the user doesn't have
+    permission to access the archive's bucket but a local copy of the archive
+    exists.
+    """
+    # Download all .wpr files.
+    if self._data:
+      return
+    if not self._bucket:
+      logging.warning('User story set in %s has no bucket specified, and '
+                      'cannot be downloaded from cloud_storage.', )
+
+    for archive_path in self._data['archives']:
+      archive_path = self._WprFileNameToPath(archive_path)
+      try:
+        cloud_storage.GetIfChanged(archive_path, self._bucket)
+      except (cloud_storage.CredentialsError, cloud_storage.PermissionError):
+        if os.path.exists(archive_path):
+          # If the archive exists, assume the user recorded their own and
+          # simply warn.
+          logging.warning('Need credentials to update WPR archive: %s',
+                          archive_path)
+        else:
+          logging.error("You either aren't authenticated or don't have "
+                        "permission to use the archives for this page set."
+                        "\nYou may need to run gsutil config.")
+          raise
 
   def WprFilePathForUserStory(self, story):
     if self.temp_target_wpr_file_path:
       return self.temp_target_wpr_file_path
     wpr_file = self._user_story_name_to_wpr_file.get(story.display_name, None)
     if wpr_file is None and isinstance(story, page_module.Page):
       # Some old pages always use the URL to identify a page rather than the
       # display_name, so try to look for that.
       wpr_file = self._user_story_name_to_wpr_file.get(story.url, None)
     if wpr_file:
@@ skipping 100 matching lines @@
 
   def _SetWprFileForUserStory(self, user_story_name, wpr_file):
     """For modifying the metadata when we're going to record a new archive."""
     old_wpr_file = self._user_story_name_to_wpr_file.get(user_story_name, None)
     if old_wpr_file:
       self._wpr_file_to_user_story_names[old_wpr_file].remove(user_story_name)
     self._user_story_name_to_wpr_file[user_story_name] = wpr_file
     if wpr_file not in self._wpr_file_to_user_story_names:
       self._wpr_file_to_user_story_names[wpr_file] = []
     self._wpr_file_to_user_story_names[wpr_file].append(user_story_name)