Tool to generate screenshots of Chrome in various languages

tools/privacy_whitepaper/screenshot_generator.py

+# Copyright 2016 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 environment
+import logging
+import os
+import pyautogui
+import time
+
+from PIL import Image
+from selenium.webdriver.common.action_chains import ActionChains
+
+SCRIPT_DEBUG = 100
+
+class ScreenshotGenerator:
+  """Base class for creation of screenshots providing utility functions."""
+
+  def __init__(self, name):
+    """Creates a new ScreenshotGenerator instance.
+
+    Args:
+      name: The name of the screenshot type.
+    """
+    self.name = name
+
+    # The Environment, if added to any.
+    self.environment = None
+    # The webdriver from the environment.
+    self.driver = None
+
+  # Mouse/Keyboard actions.
+
+  def Click(self, selector):
+    """Clicks on the element described by |selector|.
+
+    Args:
+      selector: The clicked element's CSS selector.
+    """
+
+    logging.log(SCRIPT_DEBUG, "action: Click %s" % selector)

[melandory, 2016/05/30 13:51:49]

nit:
since it's python and you're using same pattern in many places (log+wait for
element), you can write decorator:
def LogAndWait(f):
  from functools import wraps
  @wraps(f)
  def wrapper(self, selector):
    logging.log(SCRIPT_DEBUG, "action: %s %s" % (f.__name__, selector))
    element = self.WaitUntilDisplayed(selector)
    return f(self, element)
  return wrapper

And then
@LogAndWait
def Click(self, element)

+    element = self.WaitUntilDisplayed(selector)
+    element.click()
+
+  def ClickIfClickable(self, selector):
+    """Clicks on the element described by |selector| if it is clickable.
+
+    The driver's find_element_by_css_selector method defines what is clickable
+    -- anything for which it does not throw, is clickable. To be clickable,
+    the element must:
+      * exist in the DOM,
+      * be not covered by another element
+      * be inside the visible area.
+    Note that transparency does not influence clickability.
+
+    Args:
+      selector: The clicked element's CSS selector.
+
+    Returns:
+      True if the element is clickable (and was clicked on).
+      False otherwise.
+    """
+
+    logging.log(SCRIPT_DEBUG, "action: ClickIfVisible %s" % selector)
+    element = self.WaitUntilDisplayed(selector)
+    try:
+      element.click()
+      return True
+    except Exception:
+      return False
+
+  def GoTo(self, url):
+    """Navigates the main frame to |url|.
+
+    Args:
+      url: The URL of where to go to.
+    """
+
+    logging.log(SCRIPT_DEBUG, "action: GoTo %s" % self.name)
+    self.driver.get(url)
+
+  def HoverOver(self, selector):
+    """Hovers over the element described by |selector|.
+
+    Args:
+      selector: The CSS selector of the element to hover over.
+    """
+
+    logging.log(SCRIPT_DEBUG, "action: Hover %s" % selector)
+    element = self.WaitUntilDisplayed(selector)
+    hover = ActionChains(self.driver).move_to_element(element)
+    hover.perform()
+
+  # Waiting/Displaying actions.
+
+  def _ReturnElementIfDisplayed(self, selector):
+    """Returns the element described by |selector|, if displayed.
+
+    Note: This takes neither overlapping among elements nor position with
+    regards to the visible area into account.
+
+    Args:
+      selector: The CSS selector of the checked element.
+
+    Returns:
+      The element if displayed, None otherwise.
+    """
+
+    try:
+      element = self.driver.find_element_by_css_selector(selector)
+      if element:
+        logging.log(SCRIPT_DEBUG, "Found element %s" % selector)
+      else:
+        logging.log(SCRIPT_DEBUG, "Did not find element %s" % selector)
+      return element if element.is_displayed() else None
+    except Exception as e:
+      logging.log(SCRIPT_DEBUG, "Exception %s", e)
+      return None
+
+  def IsDisplayed(self, selector):
+    """Check if the element described by |selector| is displayed.
+
+    Note: This takes neither overlapping among elements nor position with
+    regards to the visible area into account.
+
+    Args:
+      selector: The CSS selector of the checked element.
+
+    Returns:
+      True if the element is in the DOM and less than 100% transparent.
+      False otherwise.
+    """
+
+    logging.log(SCRIPT_DEBUG, "action: IsDisplayed %s" % selector)
+    return self._ReturnElementIfDisplayed(selector) is not None
+
+  def Wait(self, duration):
+    """Wait for |duration| in seconds.
+
+    Args:
+      duration: The time to wait in seconds.
+
+    Raises:
+      Exception: In case the accummulated waiting limit is exceeded.
+    """
+
+    logging.log(SCRIPT_DEBUG, "action: Wait %s" % duration)
+    time.sleep(duration)
+
+  def WaitUntilDisplayed(self, selector):
+    """Waits until the element described by |selector| is displayed.
+
+    Args:
+      selector: The CSS selector of the element to wait for.
+
+    Returns:
+      The displayed element.
+    """
+
+    element = self._ReturnElementIfDisplayed(selector)
+    while not element:
+      logging.log(SCRIPT_DEBUG, "action: Waiting for %s to show up" % selector)
+      self.Wait(1)
+      element = self._ReturnElementIfDisplayed(selector)
+    return element
+
+  def CaptureContentScreenshot(self, selector):
+    """Waits until the element is displayed and captures a screenshot of it.
+
+    The screenshot is written to self.environment.output_path.
+
+    Args:
+      selector: The CSS selector of the element to take a picture of.
+    """
+
+    filename = "%s-%s.png" % (self.name, self.environment.language)
+    final_path = os.path.join(self.environment.output_path, filename)
+    tmp_path = os.path.join(self.environment.output_path, "tmp.png")
+    logging.log(SCRIPT_DEBUG, "Storing a screenshot to %s", final_path)
+
+    self.HoverOver(selector)
+    element = self.WaitUntilDisplayed(selector)
+    location, size = element.location_once_scrolled_into_view, element.size
+    x, y = location['x'], location['y']
+    w, h = size['width'], size['height']
+    self.driver.save_screenshot(tmp_path)
+    img = Image.open(tmp_path)
+    img = img.crop((x, y, x + w, y + h))
+    img.save(final_path, 'png')
+    os.remove(tmp_path)
+
+  def CaptureWindowScreenshot(self, top=0, bottom=0, left=0, right=0):
+    """Takes a screenshot of the top/bottom/left/right pixels of the browser
+    window.
+
+    This requires that the browser window is positioned on the top left of the
+    screen.
+
+    Args:
+      top: Number of pixel rows to take a screenshot of.
+      bottom: Number of pixel rows to take a screenshot of.
+      left: Number of pixel columns to take a screenshot of.
+      right: Number of pixel columns to take a screenshot of.
+    """
+    filename = "%s-%s.png" % (self.name, self.environment.language)
+    final_path = os.path.join(self.environment.output_path, filename)
+    tmp_path = os.path.join(self.environment.output_path, "tmp.png")
+
+    im = pyautogui.screenshot()
+    im.save(tmp_path)
+
+    img = Image.open(tmp_path)
+    img = img.crop((0, 0, self.driver.get_window_size()['width'],
+                    self.driver.get_window_size()['height']))
+    if top != 0:
+      img = img.crop((0, 0, img.size[0], top))
+    if bottom != 0:
+      img = img.crop((0, img.size[1] - bottom, img.size[0], img.size[1]))
+    if left != 0:
+      img = img.crop((0, 0, left, img.size[1]))
+    if right != 0:
+      img = img.crop((img.size[0] - right, 0, img.size[0], img.size[1]))
+    img.save(final_path, 'png')
+    os.remove(tmp_path)