Source code for webtraversallibrary.screenshot

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at


# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""Abstraction layer for a screenshot of a site"""

from __future__ import annotations

import io
import os
import time
from math import ceil
from pathlib import Path
from typing import Dict, List

from PIL import Image
from selenium.webdriver.remote.webdriver import WebDriver

from .color import Color
from .geometry import Point, Rectangle
from .graphics import draw_rect, draw_text
from .javascript import JavascriptWrapper

[docs]class Screenshot: """Abstraction layer for a screenshot of a site, allowing for various annotations.""" def __init__(self, name: str, image: Image.Image): = name self.image = image
[docs] @classmethod def capture_viewport(cls, name: str, driver: WebDriver, scale: float = 1.0) -> Screenshot: """ Creates a screenshot of the current viewport of a given webdriver. Scales the image by some pixel ratio, if given. Uses PIL as a backend. """ page_screenshot_png_bytes = driver.get_screenshot_as_png() image ="RGB") if scale != 1.0: new_size = [int(x * scale) for x in image.size] image = image.resize(new_size) return Screenshot(name, image)
[docs] @classmethod def capture(cls, name: str, driver: WebDriver, scale: float = 1.0, max_page_height: int = 0) -> Screenshot: """ Creates a snapshot on the given webdriver under certain conditions. """ js = JavascriptWrapper(driver) viewport = js.find_viewport() if max_page_height < viewport.height: return cls.capture_viewport(name, driver, scale) # Capture at least one viewport's height, at most the page's height minimum_length = max(int(viewport.height), js.get_full_height()) page_height = min(max_page_height, minimum_length) num_shots = ceil(page_height / viewport.height) viewport_screenshots: List[Image.Image] = [] hidden_elements: Dict[int, str] = js.hide_position_fixed_elements() try: # We start from bottom, and go up! for shot_idx in range(num_shots, -1, -1): # Don't hide elements at the top if shot_idx == 0 or viewport.height > page_height: js.show_position_fixed_elements(hidden_elements) # Scroll, hide sticky elements, and take viewport screenshot hidden = Screenshot._scroll_to(js, shot_idx * int(viewport.height)) hidden_elements.update(hidden) viewport_screenshot = cls.capture_viewport("temp", driver, scale).image if shot_idx * viewport.height > page_height - viewport.height: new_viewport_area = ( 0, 0, viewport.width, page_height - shot_idx * viewport.height, ) viewport_screenshot = viewport_screenshot.crop(new_viewport_area) # Append screenshot of the current viewport at the top of the image viewport_screenshots.insert(0, viewport_screenshot) finally: # Unhide everything and scroll back to the top js.show_position_fixed_elements(hidden_elements) total_height = sum(image.height for image in viewport_screenshots) final_screenshot ="RGB", (viewport_screenshots[0].width, total_height)) offset = 0 for viewport_screenshot in viewport_screenshots: final_screenshot.paste(viewport_screenshot, (0, offset)) offset += viewport_screenshot.height return Screenshot(name, final_screenshot)
@classmethod def load(cls, name: str, path: Path) -> Screenshot: return cls(name,
[docs] def save(self, path: Path, suffix: str = ""): """Saves screenshot to given path. Filename consists of the screenshot name and an optional suffix.""" os.makedirs(path, exist_ok=True) filename = f"{}_{suffix}.png" if suffix else f"{}.png" / filename))
def copy(self, new_name: str) -> Screenshot: return Screenshot(new_name, self.image.copy())
[docs] def highlight(self, rect: Rectangle, color: Color, text: str = "", width: int = 1): """ Draws a colored rectangle on the screenshot. Can also annotate with a text below the rectangle, if given. """ draw_rect(self.image, rect, color, width) if text: offset = 1 self.annotate(Point(rect.x + offset, rect.y + rect.height + offset), color, 12, text)
[docs] def annotate(self, top_left: Point, color: Color, size: int, text: str): """ Writes text with a given color on the screenshot. """ draw_text(self.image, top_left, color, size, text)
@property def size(self) -> Point: """Returns a (width, height) Point of the screenshot size in pixels""" return Point(self.width, self.height) @property def height(self) -> int: return self.image.height @property def width(self) -> int: return self.image.width @staticmethod def _scroll_to(js: JavascriptWrapper, y_pos: int, elements: List[str] = None): waiting_time = 0.1 hidden: Dict[int, str] = {} js.scroll_to(0, y_pos) time.sleep(waiting_time) hidden = js.hide_position_fixed_elements(elements=elements if elements else []) time.sleep(waiting_time) return hidden