status-django

  1"""
  2Headless Chromium wrapper for generating screenshots and PDFs from URLs or HTML.
  3
  4Shells out to the system chromium binary (no Playwright, no Selenium, no bundled
  5browser). Works on Alpine, Ubuntu, and macOS as long as a chromium-family binary
  6is discoverable on PATH.
  7"""
  8from __future__ import annotations
  9
 10import os
 11import shutil
 12import subprocess
 13import tempfile
 14from contextlib import contextmanager
 15from pathlib import Path
 16from typing import Iterator, Optional, Tuple
 17
 18from django.core.files import File
 19from django.core.files.storage import default_storage
 20
 21DEFAULT_VIEWPORT: Tuple[int, int] = (1280, 720)
 22DEFAULT_VIRTUAL_TIME_BUDGET_MS = 5_000
 23DEFAULT_SUBPROCESS_TIMEOUT_S = 60
 24
 25BASE_FLAGS = (
 26    "--headless=new",
 27    "--no-sandbox",
 28    "--no-zygote",
 29    "--disable-gpu",
 30    "--disable-dev-shm-usage",
 31    "--disable-software-rasterizer",
 32    "--disable-extensions",
 33    "--disable-background-networking",
 34    "--disable-crash-reporter",
 35    "--disable-logging",
 36    "--hide-scrollbars",
 37)
 38
 39
 40def _find_chromium() -> Optional[str]:
 41    for binary in ("chromium", "chromium-browser", "google-chrome"):
 42        path = shutil.which(binary)
 43        if path:
 44            return path
 45    return None
 46
 47
 48CHROMIUM_BINARY = _find_chromium()
 49
 50
 51class ChromiumError(RuntimeError):
 52    pass
 53
 54
 55@contextmanager
 56def _tempfile(suffix: str) -> Iterator[Path]:
 57    fd, raw = tempfile.mkstemp(suffix=suffix, dir="/tmp")
 58    os.close(fd)
 59    path = Path(raw)
 60    try:
 61        yield path
 62    finally:
 63        path.unlink(missing_ok=True)
 64
 65
 66@contextmanager
 67def _html_tempfile(html: str) -> Iterator[str]:
 68    fd, raw = tempfile.mkstemp(suffix=".html", dir="/tmp")
 69    path = Path(raw)
 70    try:
 71        with os.fdopen(fd, "w", encoding="utf-8") as fp:
 72            fp.write(html)
 73        yield f"file://{path}"
 74    finally:
 75        path.unlink(missing_ok=True)
 76
 77
 78def _run(args: list[str], timeout: int) -> None:
 79    if not CHROMIUM_BINARY:
 80        raise ChromiumError(
 81            "No chromium binary found on PATH (tried: chromium, "
 82            "chromium-browser, google-chrome)"
 83        )
 84    cmd = [CHROMIUM_BINARY, *BASE_FLAGS, *args]
 85    try:
 86        subprocess.run(cmd, check=True, capture_output=True, timeout=timeout)
 87    except subprocess.TimeoutExpired as exc:
 88        raise ChromiumError(f"chromium timed out after {timeout}s") from exc
 89    except subprocess.CalledProcessError as exc:
 90        stderr = (exc.stderr or b"").decode("utf-8", errors="replace").strip()
 91        raise ChromiumError(
 92            f"chromium exited {exc.returncode}: {stderr or '(no stderr)'}"
 93        ) from exc
 94
 95
 96def _save(source: Path, filename: str) -> str:
 97    if default_storage.exists(filename):
 98        default_storage.delete(filename)
 99    with source.open("rb") as fp:
100        default_storage.save(filename, File(fp))
101    return default_storage.url(filename)
102
103
104def generate_screenshot_from_url(
105    url: str,
106    filename: str,
107    *,
108    viewport: Tuple[int, int] = DEFAULT_VIEWPORT,
109    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
110    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
111) -> str:
112    with _tempfile(".png") as out:
113        _run(
114            [
115                f"--screenshot={out}",
116                f"--window-size={viewport[0]},{viewport[1]}",
117                f"--virtual-time-budget={virtual_time_budget_ms}",
118                url,
119            ],
120            timeout,
121        )
122        return _save(out, filename)
123
124
125def generate_screenshot_from_html(
126    html: str,
127    filename: str,
128    *,
129    viewport: Tuple[int, int] = DEFAULT_VIEWPORT,
130    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
131    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
132) -> str:
133    with _html_tempfile(html) as url:
134        return generate_screenshot_from_url(
135            url,
136            filename,
137            viewport=viewport,
138            virtual_time_budget_ms=virtual_time_budget_ms,
139            timeout=timeout,
140        )
141
142
143def generate_pdf_from_url(
144    url: str,
145    filename: str,
146    *,
147    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
148    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
149) -> str:
150    with _tempfile(".pdf") as out:
151        _run(
152            [
153                f"--print-to-pdf={out}",
154                "--no-pdf-header-footer",
155                f"--virtual-time-budget={virtual_time_budget_ms}",
156                url,
157            ],
158            timeout,
159        )
160        return _save(out, filename)
161
162
163def generate_pdf_from_html(
164    html: str,
165    filename: str,
166    *,
167    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
168    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
169) -> str:
170    with _html_tempfile(html) as url:
171        return generate_pdf_from_url(
172            url,
173            filename,
174            virtual_time_budget_ms=virtual_time_budget_ms,
175            timeout=timeout,
176        )