[Bugfix & Enhancement] Screenshot

splinter/driver/webdriver/__init__.py

@@ -10,6 +10,7 @@
 import sys
 import tempfile
 import time
+import base64
 from contextlib import contextmanager
 import warnings
 
@@ -40,6 +41,19 @@ def alert_enter(self):
 def alert_exit(self, type, value, traceback):
     pass
 
+def create_file_from_base64_string(base64_string, filename):
+    file = base64.b64decode(base64_string.encode('ascii'))
+
+    try:
+        with open(filename, 'wb') as f:
+            f.write(file)
+    except IOError:
+        return False
+    finally:
+        del file
+
+    return True
+
 Alert.__enter__ = alert_enter
 Alert.__exit__ = alert_exit
 Alert.fill_with = Alert.send_keys
@@ -600,19 +614,93 @@ def uncheck(self, name):
         self.find_by_name(name).first.uncheck()
 
     def screenshot(self, name="", suffix=".png", full=False):
+        warnings.warn('Deprecated, use `self.capture_screenshot()` instead.')
 
         name = name or ""
 
         (fd, filename) = tempfile.mkstemp(prefix=name, suffix=suffix)
         # don't hold the file
         os.close(fd)
 
+        self.capture_screenshot(filename, full)
+
+        return filename
+
+    def capture_viewport(self, viewport, filename, waiting_time=0):
+        base64_string = self.capture_viewport_as_base64(viewport, waiting_time)
+
+        return create_file_from_base64_string(base64_string, filename)
+
+
+    def capture_viewport_as_base64(self, viewport, waiting_time=0):
+        self.full_screen()
+
+        # trigger the `scroll` event to ensure lazy-load images can be rendered.
+        self.execute_script("window.dispatchEvent(new Event('scroll'))")
+
+        if waiting_time > 0: time.sleep(waiting_time)
+
+        response = self.driver.execute_cdp_cmd('Page.captureScreenshot', {
+            'clip': {
+                'x': viewport['x'],
+                'y': viewport['y'],
+                'width': viewport['width'],
+                'height': viewport['height'],
+                'scale': 1
+            }
+        })
+
+        self.recover_screen()
+
+        return response['data']
+
+
+    def capture_screenshot(self, filename, full=False, viewport=None, waiting_time=0):
         if full:
             self.full_screen()
 
-        self.driver.get_screenshot_as_file(filename)
-        self.recover_screen()
-        return filename
+            # trigger the `scroll` event to ensure lazy-load images can be rendered.
+            self.execute_script("window.dispatchEvent(new Event('scroll'))")
+
+            if waiting_time > 0: time.sleep(waiting_time)
+
+            result = self.driver.get_screenshot_as_file(filename)
+
+            self.recover_screen()
+
+            return result
+
+        elif viewport:
+            return self.capture_viewport(viewport, filename, waiting_time)
+
+        else:
+            if waiting_time > 0: time.sleep(waiting_time)
+
+            return self.driver.get_screenshot_as_file(filename)
+
+
+    def capture_screenshot_as_base64(self, full=False, viewport=None, waiting_time=0):
+        if full:
+            self.full_screen()
+
+            # trigger the `scroll` event to ensure lazy-load images can be rendered.
+            self.execute_script("window.dispatchEvent(new Event('scroll'))")
+
+            if waiting_time > 0: time.sleep(waiting_time)
+
+            result = self.driver.get_screenshot_as_base64()
+
+            self.recover_screen()
+
+            return result
+
+        elif viewport:
+            return self.capture_viewport_as_base64(viewport, waiting_time)
+
+        else:
+            if waiting_time > 0: time.sleep(waiting_time)
+
+            return self.driver.get_screenshot_as_base64()
 
     def select(self, name, value):
         self.find_by_xpath(
@@ -631,18 +719,33 @@ def quit(self):
             pass
 
     def full_screen(self):
-        self.ori_window_size = self.driver.get_window_size()
-        width = self.driver.execute_script("return Math.max(document.body.scrollWidth, document.body.offsetWidth);")
-        height = self.driver.execute_script("return Math.max(document.body.scrollHeight, document.body.offsetHeight);")
-        self.driver.set_window_size(width, height)
+        self.set_original_window_size()
+        width = self.driver.execute_script("return Math.max(document.body.scrollWidth, document.body.offsetWidth, document.documentElement.scrollWidth);")
+        height = self.driver.execute_script("return Math.max(document.body.scrollHeight, document.body.offsetHeight, document.documentElement.scrollHeight);")
+        self.set_viewport_size(width, height)
 
     def recover_screen(self):
         if self.ori_window_size:
             width = self.ori_window_size.get('width')
             height = self.ori_window_size.get('height')
-            self.driver.set_window_size(width, height)
+            self.set_viewport_size(width, height)
             self.ori_window_size = None
 
+    def set_original_window_size(self):
+        self.ori_window_size = self.driver.get_window_size()
+
+    def set_viewport_size(self, width, height):
+        device_metrics = {
+            'width': width,
+            'height': height,
+            'deviceScaleFactor': self.execute_script('return window.devicePixelRatio'),
+            'mobile': self.execute_script("return typeof window.orientation !== 'undefined'")
+        }
+
+        self.driver.execute_cdp_cmd('Emulation.setDeviceMetricsOverride', device_metrics)
+
+        return self
+
     def html_snapshot(self, name="", suffix=".html", encoding='utf-8'):
         """Write the current html to a file."""
         name = name or ""
@@ -873,52 +976,51 @@ def drag_and_drop(self, droppable):
         ActionChains(self.parent.driver).drag_and_drop(self._element, droppable._element).perform()
 
     def screenshot(self, name='', suffix='.png', full=False):
+        warnings.warn('Deprecated, use `self.capture_screenshot()` instead.')
+
         name = name or ''
 
         (fd, filename) = tempfile.mkstemp(prefix=name, suffix=suffix)
         # don't hold the file
         os.close(fd)
 
-        if full:
-            self.parent.full_screen()
-        target = self.screenshot_as_png()
-        self.parent.recover_screen()
-        target.save(filename)
+        self.capture_screenshot(filename)
 
         return filename
 
-    def screenshot_as_png(self):
-        try:
-            from PIL import Image
-        except ImportError:
-            raise NotImplementedError('Element screenshot need the Pillow dependency. '
-                                      'Please use "pip install Pillow" install it.')
+    # TODO Uncomment this method and remove the hotfix one once the upstream issue(https://bugs.chromium.org/p/chromedriver/issues/detail?id=3154) has been fixed.
+    # def capture_screenshot(self, filename, waiting_time=0):
+    #     if waiting_time > 0:
+    #         time.sleep(waiting_time)
 
-        full_screen_png = self.parent.driver.get_screenshot_as_png()
+    #     return self._element.screenshot(filename)
 
-        full_screen_bytes = BytesIO(full_screen_png)
+    def capture_screenshot(self, filename, waiting_time=0):
+        base64_string = self.capture_screenshot_as_base64(waiting_time)
 
-        im = Image.open(full_screen_bytes)
-        im_width, im_height = im.size[0], im.size[1]
-        window_size = self.parent.driver.get_window_size()
-        window_width = window_size['width']
+        return create_file_from_base64_string(base64_string, filename)
 
-        ratio = im_width * 1.0 / window_width
-        height_ratio = im_height / ratio
 
-        im = im.resize((int(window_width), int(height_ratio)))
+    # TODO Uncomment this method and remove the hotfix one once the upstream issue(https://bugs.chromium.org/p/chromedriver/issues/detail?id=3154) has been fixed.
+    # def capture_screenshot_as_base64(self, waiting_time=0):
+    #     if waiting_time > 0:
+    #         time.sleep(waiting_time)
+
+    #     return self._element.screenshot_as_base64
 
+    def capture_screenshot_as_base64(self, waiting_time=0):
         location = self._element.location
-        x, y = location['x'], location['y']
+        size = self._element.size
 
-        pic_size = self._element.size
-        w, h = pic_size['width'], pic_size['height']
+        viewport = {
+            'x': location['x'],
+            'y': location['y'],
+            'width': size['width'],
+            'height': size['height']
+        }
 
-        box = x, y, x + w, y + h
-        box = [int(i) for i in box]
-        target = im.crop(box)
+        return self.parent.capture_viewport_as_base64(viewport, waiting_time)
 
-        return target
 
     def __getitem__(self, attr):
         return self._element.get_attribute(attr)