rebaseline_server: add ImagePair class, a step towards new intermediate JSON schema

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.
 """
 
 import contextlib
 import logging
 import os
+import re
 import shutil
 import urllib
 try:
   from PIL import Image, ImageChops
 except ImportError:
   raise ImportError('Requires PIL to be installed; see '
                     + 'http://www.pythonware.com/products/pil/')
 
 DEFAULT_IMAGE_SUFFIX = '.png'
 DEFAULT_IMAGES_SUBDIR = 'images'
 
+DISALLOWED_FILEPATH_CHAR_REGEX = re.compile('[^\w\-]')
+
 DIFFS_SUBDIR = 'diffs'
 WHITEDIFFS_SUBDIR = 'whitediffs'
 
 VALUES_PER_BAND = 256
 
 
 class DiffRecord(object):
   """ Record of differences between two images. """
 
   def __init__(self, storage_root,
@@ skipping 18 matching lines @@
           guarantee uniqueness)
       actual_image_url: file 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 makes _download_and_open_image raise
     # an exception if images are not found locally (instead of trying to
     # download them).
     expected_image = _download_and_open_image(
         os.path.join(storage_root, expected_images_subdir,
                      str(expected_image_locator) + image_suffix),
         expected_image_url)
     actual_image = _download_and_open_image(
         os.path.join(storage_root, actual_images_subdir,
@@ skipping 51 matching lines @@
   def get_weighted_diff_measure(self):
     """Returns a weighted measure of image diffs, as a float between 0 and 100
     (inclusive)."""
     return self._weighted_diff_measure
 
   def get_max_diff_per_channel(self):
     """Returns the maximum difference between the expected and actual images
     for each R/G/B channel, as a list."""
     return self._max_diff_per_channel
 
+  def as_dict(self):
+    """Returns a dictionary representation of this DiffRecord, as needed when
+    constructing the JSON representation."""
+    return {
+        'numDifferingPixels': self._num_pixels_differing,
+        'percentDifferingPixels': self.get_percent_pixels_differing(),
+        'weightedDiffMeasure': self.get_weighted_diff_measure(),
+        'maxDiffPerChannel': self._max_diff_per_channel,
+    }
+
 
 class ImageDiffDB(object):
   """ Calculates differences between image pairs, maintaining a database of
   them for download."""
 
   def __init__(self, storage_root):
     """
     Args:
       storage_root: string; root path within the DB will store all of its stuff
     """
@@ skipping 22 matching lines @@
           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
       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,
             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:
         logging.exception('got exception while creating new DiffRecord')
         return
       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 = (expected_image_locator, actual_image_locator)
+    key = (_sanitize_locator(expected_image_locator),
+           _sanitize_locator(actual_image_locator))
     return self._diff_dict[key]
 
 
 # Utility functions
 
 def _calculate_weighted_diff_metric(histogram, num_pixels):
   """Given the histogram of a diff image (per-channel diff at each
   pixel between two images), calculate the weighted diff metric (a
   stab at how different the two images really are).
 
@@ skipping 108 matching lines @@
 
 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)
 
+def _sanitize_locator(locator):
+  """Returns a sanitized version of a locator (one in which we know none of the
+  characters will have special meaning in filenames).
+
+  Args:
+    locator: string, or something that can be represented as a string
+  """
+  return DISALLOWED_FILEPATH_CHAR_REGEX.sub('_', str(locator))
+
 def _get_difference_locator(expected_image_locator, actual_image_locator):
   """Returns the locator string used to look up the diffs between expected_image
   and actual_image.
 
   Args:
     expected_image_locator: locator string pointing at expected image
     actual_image_locator: locator string pointing at actual image
 
-  Returns: locator where the diffs between expected and actual images can be
-      found
+  Returns: already-sanitized locator where the diffs between expected and
+      actual images can be found
   """
-  return "%s-vs-%s" % (expected_image_locator, actual_image_locator)
+  return "%s-vs-%s" % (_sanitize_locator(expected_image_locator),
+                       _sanitize_locator(actual_image_locator))