screenshot.py

import argparse
import io
from pathlib import Path
import time

from PIL import Image
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from webdriver_manager.chrome import ChromeDriverManager


def parse_args():
    parser = argparse.ArgumentParser(
        description="Capture a screenshot of a URL using Selenium."
    )
    parser.add_argument(
        "url",
        nargs="?",
        default="https://example.com",
        help="Page URL to capture.",
    )
    parser.add_argument(
        "-o",
        "--output",
        default="screenshots/screenshot.png",
        help="Output file path.",
    )
    output_group = parser.add_mutually_exclusive_group()
    output_group.add_argument(
        "--element",
        help="CSS selector for an element screenshot.",
    )
    output_group.add_argument(
        "--full-page",
        action="store_true",
        help="Scroll and stitch a full-page screenshot.",
    )
    parser.add_argument(
        "--headed",
        action="store_true",
        help="Run with a visible browser window.",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=1920,
        help="Viewport width.",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=1080,
        help="Viewport height.",
    )
    parser.add_argument(
        "--extra-wait",
        type=float,
        default=0.0,
        help="Extra seconds to wait after the body loads.",
    )
    parser.add_argument(
        "--scroll-wait",
        type=float,
        default=0.2,
        help="Seconds to wait after each scroll when stitching.",
    )
    return parser.parse_args()


def build_driver(headless, width, height):
    options = webdriver.ChromeOptions()
    if headless:
        options.add_argument("--headless")
    options.add_argument(f"--window-size={width},{height}")
    options.add_argument("--disable-gpu")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")
    return webdriver.Chrome(
        service=ChromeService(ChromeDriverManager().install()),
        options=options,
    )


def full_page_screenshot(driver, output_path, scroll_wait):
    metrics = driver.execute_script(
        """
        const body = document.body;
        const html = document.documentElement;
        const totalHeight = Math.max(
            body.scrollHeight, body.offsetHeight, body.clientHeight,
            html.scrollHeight, html.offsetHeight, html.clientHeight
        );
        const viewportHeight = window.innerHeight;
        return { totalHeight, viewportHeight };
        """
    )

    total_height = int(metrics["totalHeight"])
    viewport_height = int(metrics["viewportHeight"])

    if viewport_height <= 0:
        raise RuntimeError("Viewport height is zero")

    last_scroll = max(total_height - viewport_height, 0)
    positions = list(range(0, total_height, viewport_height))
    if not positions or positions[-1] != last_scroll:
        positions.append(last_scroll)

    screenshots = []

    for offset in positions:
        driver.execute_script(
            "document.documentElement.scrollTo(0, arguments[0]);"
            "document.body.scrollTo(0, arguments[0]);",
            offset,
        )
        time.sleep(scroll_wait)

        png = driver.get_screenshot_as_png()
        screenshots.append(Image.open(io.BytesIO(png)))

    if not screenshots:
        raise RuntimeError("No screenshots captured")

    scale = screenshots[0].height / viewport_height
    total_height_px = int(total_height * scale)
    total_width_px = screenshots[0].width

    stitched = Image.new("RGB", (total_width_px, total_height_px))
    y = 0

    for img in screenshots:
        remaining = total_height_px - y
        if remaining <= 0:
            break
        if img.height > remaining:
            img = img.crop((0, 0, img.width, remaining))
        stitched.paste(img, (0, y))
        y += img.height

    stitched.save(output_path)


def main():
    args = parse_args()
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    driver = build_driver(not args.headed, args.width, args.height)
    try:
        driver.get(args.url)
        WebDriverWait(driver, 10).until(
            EC.presence_of_element_located((By.TAG_NAME, "body"))
        )
        if args.extra_wait > 0:
            time.sleep(args.extra_wait)

        if args.full_page:
            full_page_screenshot(driver, str(output_path), args.scroll_wait)
        elif args.element:
            element = driver.find_element(By.CSS_SELECTOR, args.element)
            element.screenshot(str(output_path))
        else:
            driver.save_screenshot(str(output_path))
    finally:
        driver.quit()

    print(f"Saved screenshot to {output_path}")


if __name__ == "__main__":
    main()