teach rebaseline_server to generate diffs of rendered SKPs

gm/rebaseline_server/imagediffdb.py

 #!/usr/bin/python
 
 """
 Copyright 2013 Google Inc.
 
 Use of this source code is governed by a BSD-style license that can be
 found in the LICENSE file.
 
 Calulate differences between image pairs, and store them in a database.
 """
 
 # System-level imports
 import contextlib
 import json
 import logging
 import os
 import re
 import shutil
 import tempfile
 import urllib
 
 # Must fix up PYTHONPATH before importing from within Skia
 import fix_pythonpath  # pylint: disable=W0611
 
 # Imports from within Skia
 import find_run_binary
+from py.utils import gs_utils
+
 
 SKPDIFF_BINARY = find_run_binary.find_path_to_program('skpdiff')
 
 DEFAULT_IMAGE_SUFFIX = '.png'
 DEFAULT_IMAGES_SUBDIR = 'images'
 
 DISALLOWED_FILEPATH_CHAR_REGEX = re.compile('[^\w\-]')
 
 RGBDIFFS_SUBDIR = 'diffs'
 WHITEDIFFS_SUBDIR = 'whitediffs'
 
 # Keys used within DiffRecord dictionary representations.
 # NOTE: Keep these in sync with static/constants.js
 KEY__DIFFERENCES__MAX_DIFF_PER_CHANNEL = 'maxDiffPerChannel'
 KEY__DIFFERENCES__NUM_DIFF_PIXELS = 'numDifferingPixels'
 KEY__DIFFERENCES__PERCENT_DIFF_PIXELS = 'percentDifferingPixels'
 KEY__DIFFERENCES__PERCEPTUAL_DIFF = 'perceptualDifference'
 
 
+class GSObject(object):
+  """Reference to a file, dir, or entire bucket in Google Storage."""
+  # EPOGER: move this class into gs_utils.py ?
+  def __init__(self, gs_url=None, bucket=None, path=None):
+    """
+    Create a reference to a GS object, using EITHER gs_url OR bucket/path.
+
+    Args:
+      gs_url: URL pointing to this file, e.g. "gs://bucketname/dir/file".
+          If this is set, the other params must NOT be set.
+      bucket: name of the bucket, e.g. "bucketname".
+          If this is set, the "gs_url" param must NOT be set.
+      path: Posix-style path within the bucket, e.g. "dir/file".
+          If this is set, the "gs_url" param must NOT be set.
+    """
+    if gs_url != None:
+      if bucket != None or path != None:
+        raise Exception(
+            'gs_url is set, so neither bucket nor path should be set')
+      (self.bucket, self.path) = gs_utils.GSUtils.split_gs_url(gs_url)
+    else:
+      if not bucket:
+        raise Exception('gs_url is not set, so bucket must be set')
+      self.bucket = bucket
+      self.path = path or ''
+
+  def as_gs_url(self):
+    """Returns a gs:// URL pointing at this object in Google Storage."""
+    url = 'gs://' + self.bucket
+    if self.path:
+      url += '/' + self.path
+    return url
+
+  def as_http_url(self):
+    """Returns an http:// URL pointing at this object in Google Storage."""
+    url = 'http://%s.commondatastorage.googleapis.com' % self.bucket
+    if self.path:
+      url += '/' + self.path
+    return url
+
+
 class DiffRecord(object):
   """ Record of differences between two images. """
 
-  def __init__(self, storage_root,
+  def __init__(self, gs, storage_root,
                expected_image_url, expected_image_locator,
                actual_image_url, actual_image_locator,
                expected_images_subdir=DEFAULT_IMAGES_SUBDIR,
                actual_images_subdir=DEFAULT_IMAGES_SUBDIR,
                image_suffix=DEFAULT_IMAGE_SUFFIX):
     """Download this pair of images (unless we already have them on local disk),
     and prepare a DiffRecord for them.
 
     TODO(epoger): Make this asynchronously download images, rather than blocking
     until the images have been downloaded and processed.
 
     Args:
+      gs: instance of GSUtils object we can use to download images
       storage_root: root directory on local disk within which we store all
           images
-      expected_image_url: file or HTTP url from which we will download the
+      expected_image_url: file, GS, or HTTP url from which we will download the
           expected image
       expected_image_locator: a unique ID string under which we will store the
           expected image within storage_root (probably including a checksum to
           guarantee uniqueness)
-      actual_image_url: file or HTTP url from which we will download the
+      actual_image_url: file, GS, or HTTP url from which we will download the
           actual image
       actual_image_locator: a unique ID string under which we will store the
           actual image within storage_root (probably including a checksum to
           guarantee uniqueness)
       expected_images_subdir: the subdirectory expected images are stored in.
       actual_images_subdir: the subdirectory actual images are stored in.
       image_suffix: the suffix of images.
     """
     expected_image_locator = _sanitize_locator(expected_image_locator)
     actual_image_locator = _sanitize_locator(actual_image_locator)
 
     # Download the expected/actual images, if we don't have them already.
     # TODO(rmistry): Add a parameter that just tries to use already-present
     # image files rather than downloading them.
     expected_image_file = os.path.join(
         storage_root, expected_images_subdir,
         str(expected_image_locator) + image_suffix)
     actual_image_file = os.path.join(
         storage_root, actual_images_subdir,
         str(actual_image_locator) + image_suffix)
     try:
-      _download_file(expected_image_file, expected_image_url)
+      _download_file(gs, expected_image_file, expected_image_url)
     except Exception:
       logging.exception('unable to download expected_image_url %s to file %s' %
                         (expected_image_url, expected_image_file))
       raise
     try:
-      _download_file(actual_image_file, actual_image_url)
+      _download_file(gs, actual_image_file, actual_image_url)
     except Exception:
       logging.exception('unable to download actual_image_url %s to file %s' %
                         (actual_image_url, actual_image_file))
       raise
 
     # Get all diff images and values from skpdiff binary.
     skpdiff_output_dir = tempfile.mkdtemp()
     try:
       skpdiff_summary_file = os.path.join(skpdiff_output_dir,
                                           'skpdiff-output.json')
@@ skipping 96 matching lines @@
             self.get_percent_pixels_differing(),
         KEY__DIFFERENCES__MAX_DIFF_PER_CHANNEL: self._max_diff_per_channel,
         KEY__DIFFERENCES__PERCEPTUAL_DIFF: self._perceptual_difference,
     }
 
 
 class ImageDiffDB(object):
   """ Calculates differences between image pairs, maintaining a database of
   them for download."""
 
-  def __init__(self, storage_root):
+  def __init__(self, storage_root, gs=None):
     """
     Args:
       storage_root: string; root path within the DB will store all of its stuff
+      gs: instance of GSUtils object we can use to download images
     """
     self._storage_root = storage_root
+    self._gs = gs
 
     # Dictionary of DiffRecords, keyed by (expected_image_locator,
     # actual_image_locator) tuples.
     self._diff_dict = {}
 
   @property
   def storage_root(self):
     return self._storage_root
 
   def add_image_pair(self,
                      expected_image_url, expected_image_locator,
                      actual_image_url, actual_image_locator):
     """Download this pair of images (unless we already have them on local disk),
     and prepare a DiffRecord for them.
 
     TODO(epoger): Make this asynchronously download images, rather than blocking
     until the images have been downloaded and processed.
     When we do that, we should probably add a new method that will block
     until all of the images have been downloaded and processed.  Otherwise,
     we won't know when it's safe to start calling get_diff_record().
     jcgregorio notes: maybe just make ImageDiffDB thread-safe and create a
     thread-pool/worker queue at a higher level that just uses ImageDiffDB?
 
     Args:
-      expected_image_url: file or HTTP url from which we will download the
+      expected_image_url: file, GS, or HTTP url from which we will download the
           expected image
       expected_image_locator: a unique ID string under which we will store the
           expected image within storage_root (probably including a checksum to
           guarantee uniqueness)
-      actual_image_url: file or HTTP url from which we will download the
+      actual_image_url: file, GS, or HTTP url from which we will download the
           actual image
       actual_image_locator: a unique ID string under which we will store the
           actual image within storage_root (probably including a checksum to
           guarantee uniqueness)
     """
     expected_image_locator = _sanitize_locator(expected_image_locator)
     actual_image_locator = _sanitize_locator(actual_image_locator)
     key = (expected_image_locator, actual_image_locator)
     if not key in self._diff_dict:
       try:
         new_diff_record = DiffRecord(
-            self._storage_root,
+            self._gs, self._storage_root,
             expected_image_url=expected_image_url,
             expected_image_locator=expected_image_locator,
             actual_image_url=actual_image_url,
             actual_image_locator=actual_image_locator)
       except Exception:
         # If we can't create a real DiffRecord for this (expected, actual) pair,
         # store None and the UI will show whatever information we DO have.
         # Fixes http://skbug.com/2368 .
         logging.exception(
             'got exception while creating a DiffRecord for '
             'expected_image_url=%s , actual_image_url=%s; returning None' % (
                 expected_image_url, actual_image_url))
         new_diff_record = None
       self._diff_dict[key] = new_diff_record
 
   def get_diff_record(self, expected_image_locator, actual_image_locator):
     """Returns the DiffRecord for this image pair.
 
     Raises a KeyError if we don't have a DiffRecord for this image pair.
     """
     key = (_sanitize_locator(expected_image_locator),
            _sanitize_locator(actual_image_locator))
     return self._diff_dict[key]
 
 
 # Utility functions
 
-def _download_file(local_filepath, url):
+def _download_file(gs, local_filepath, url):
   """Download a file from url to local_filepath, unless it is already there.
 
   Args:
+    gs: instance of GSUtils object, in case the url points at Google Storage
     local_filepath: path on local disk where the image should be stored
-    url: URL from which we can download the image if we don't have it yet
+    url: HTTP or GS URL from which we can download the image if we don't have
+        it yet
   """
   if not os.path.exists(local_filepath):
     _mkdir_unless_exists(os.path.dirname(local_filepath))
-    with contextlib.closing(urllib.urlopen(url)) as url_handle:
-      with open(local_filepath, 'wb') as file_handle:
-        shutil.copyfileobj(fsrc=url_handle, fdst=file_handle)
+    if gs_utils.GSUtils.is_gs_url(url):
+      (bucket, path) = gs_utils.GSUtils.split_gs_url(url)
+      gs.download_file(source_bucket=bucket, source_path=path,
+                       dest_path=local_filepath)
+    else:
+      with contextlib.closing(urllib.urlopen(url)) as url_handle:
+        with open(local_filepath, 'wb') as file_handle:
+          shutil.copyfileobj(fsrc=url_handle, fdst=file_handle)
 
 
 def _mkdir_unless_exists(path):
   """Unless path refers to an already-existing directory, create it.
 
   Args:
     path: path on local disk
   """
   if not os.path.isdir(path):
     os.makedirs(path)
@@ skipping 18 matching lines @@
 
   Args:
     expected_image_locator: locator string pointing at expected image
     actual_image_locator: locator string pointing at actual image
 
   Returns: already-sanitized locator where the diffs between expected and
       actual images can be found
   """
   return "%s-vs-%s" % (_sanitize_locator(expected_image_locator),
                        _sanitize_locator(actual_image_locator))