feat: full-page screenshots (firefox raw command, chromium PIL stitch)

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

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