From ab50a984a8a656cebe9520cf4da06bb48f93cb23 Mon Sep 17 00:00:00 2001 From: Mathew Sir Guest the best Date: Tue, 7 Jul 2026 23:14:53 -0600 Subject: [PATCH] feat: full-page screenshots (firefox raw command, chromium PIL stitch) --- src/wabot/screenshot.py | 73 ++++++++++++++++++++++++++++++++++- tests/unit/test_screenshot.py | 68 ++++++++++++++++++++++++++++++++ 2 files changed, 140 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_screenshot.py diff --git a/src/wabot/screenshot.py b/src/wabot/screenshot.py index e7de733..6dbdff7 100644 --- a/src/wabot/screenshot.py +++ b/src/wabot/screenshot.py @@ -1 +1,72 @@ -"""Full-page screenshots. Implemented in the screenshot task.""" +"""Full-page screenshots for both browser families. + +Firefox: geckodriver has a native endpoint, but selenium only exposes it +on ``webdriver.Firefox`` — not on ``Remote`` — so the raw command is +registered by hand and works for both. + +Chromium: no full-page endpoint exists; the page is captured viewport by +viewport while scrolling, then stitched with PIL (port of the legacy +workaround, using in-memory PNGs instead of temp files). +""" + +from __future__ import annotations + +import base64 +import logging +import time +from io import BytesIO + +from PIL import Image + +LOGGER = logging.getLogger("wabot") + +FULL_PAGE_SCREENSHOT = "fullPageScreenshot" +_SCROLL_SETTLE_SECONDS = 0.2 + + +def save_full_page(driver, filename: str) -> bool: + """Save a full-page PNG of the current page to ``filename``.""" + if driver.capabilities.get("browserName") == "firefox": + return _firefox_full_page(driver, filename) + return _stitch_chromium(driver, filename) + + +def _firefox_full_page(driver, filename: str) -> bool: + driver.command_executor._commands.setdefault( + FULL_PAGE_SCREENSHOT, ("GET", "/session/$sessionId/moz/screenshot/full") + ) + b64 = driver.execute(FULL_PAGE_SCREENSHOT)["value"] + with open(filename, "wb") as fp: + fp.write(base64.b64decode(b64)) + return True + + +def _stitch_chromium(driver, filename: str) -> bool: + total_width = driver.execute_script("return document.body.parentNode.scrollWidth") + total_height = driver.execute_script("return document.body.parentNode.scrollHeight") + viewport_width = driver.execute_script("return document.documentElement.clientWidth") + viewport_height = driver.execute_script("return window.innerHeight") + + if total_width <= viewport_width and total_height <= viewport_height: + with open(filename, "wb") as fp: + fp.write(driver.get_screenshot_as_png()) + return True + + stitched = Image.new("RGB", (total_width, total_height)) + y = 0 + while y < total_height: + x = 0 + while x < total_width: + driver.execute_script(f"window.scrollTo({x}, {y})") + time.sleep(_SCROLL_SETTLE_SECONDS) + tile = Image.open(BytesIO(driver.get_screenshot_as_png())) + # the last row/column can't scroll a full viewport: paste aligned + # to the bottom/right edge instead of duplicating content + paste_x = min(x, total_width - viewport_width) + paste_y = min(y, total_height - viewport_height) + stitched.paste(tile, (max(paste_x, 0), max(paste_y, 0))) + x += viewport_width + y += viewport_height + stitched.save(filename) + LOGGER.debug("stitched %sx%s screenshot -> %s", total_width, total_height, filename) + return True diff --git a/tests/unit/test_screenshot.py b/tests/unit/test_screenshot.py new file mode 100644 index 0000000..617f129 --- /dev/null +++ b/tests/unit/test_screenshot.py @@ -0,0 +1,68 @@ +import base64 +from io import BytesIO +from unittest.mock import MagicMock + +from PIL import Image + +from wabot.screenshot import FULL_PAGE_SCREENSHOT, save_full_page + + +def png_bytes(width, height, color=(200, 30, 30)): + buf = BytesIO() + Image.new("RGB", (width, height), color).save(buf, format="PNG") + return buf.getvalue() + + +class TestFirefoxPath: + def test_registers_raw_command_and_decodes_base64(self, tmp_path): + driver = MagicMock(name="driver") + driver.capabilities = {"browserName": "firefox"} + driver.command_executor._commands = {} + driver.execute.return_value = {"value": base64.b64encode(png_bytes(50, 80)).decode()} + out = tmp_path / "shot.png" + assert save_full_page(driver, str(out)) is True + # helper methods only exist on webdriver.Firefox, so the raw + # geckodriver command must have been registered for Remote support + assert driver.command_executor._commands[FULL_PAGE_SCREENSHOT] == ( + "GET", + "/session/$sessionId/moz/screenshot/full", + ) + driver.execute.assert_called_once_with(FULL_PAGE_SCREENSHOT) + assert Image.open(out).size == (50, 80) + + +class TestChromiumPath: + def test_single_viewport_page_is_saved_directly(self, tmp_path): + driver = MagicMock(name="driver") + driver.capabilities = {"browserName": "chrome"} + # page fits in one viewport: total == viewport + driver.execute_script.side_effect = lambda script: { + "return document.body.parentNode.scrollWidth": 100, + "return document.body.parentNode.scrollHeight": 60, + "return document.documentElement.clientWidth": 100, + "return window.innerHeight": 60, + }.get(script) + driver.get_screenshot_as_png.return_value = png_bytes(100, 60) + out = tmp_path / "shot.png" + assert save_full_page(driver, str(out)) is True + assert Image.open(out).size == (100, 60) + + def test_tall_page_is_stitched_from_tiles(self, tmp_path): + driver = MagicMock(name="driver") + driver.capabilities = {"browserName": "chrome"} + driver.execute_script.side_effect = lambda script: { + "return document.body.parentNode.scrollWidth": 100, + "return document.body.parentNode.scrollHeight": 150, # 3 tiles of 60 + "return document.documentElement.clientWidth": 100, + "return window.innerHeight": 60, + }.get(script) + driver.get_screenshot_as_png.return_value = png_bytes(100, 60) + out = tmp_path / "shot.png" + assert save_full_page(driver, str(out)) is True + assert Image.open(out).size == (100, 150) + # scrolled at least twice beyond the initial position + scroll_calls = [ + c for c in driver.execute_script.call_args_list + if c.args and str(c.args[0]).startswith("window.scrollTo") + ] + assert len(scroll_calls) >= 2