analytics-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
 48class ChromiumError(RuntimeError):
 49    pass
 50
 51
 52@contextmanager
 53def _tempfile(suffix: str) -> Iterator[Path]:
 54    fd, raw = tempfile.mkstemp(suffix=suffix, dir="/tmp")
 55    os.close(fd)
 56    path = Path(raw)
 57    try:
 58        yield path
 59    finally:
 60        path.unlink(missing_ok=True)
 61
 62
 63@contextmanager
 64def _html_tempfile(html: str) -> Iterator[str]:
 65    fd, raw = tempfile.mkstemp(suffix=".html", dir="/tmp")
 66    path = Path(raw)
 67    try:
 68        with os.fdopen(fd, "w", encoding="utf-8") as fp:
 69            fp.write(html)
 70        yield f"file://{path}"
 71    finally:
 72        path.unlink(missing_ok=True)
 73
 74
 75def _run(args: list[str], timeout: int) -> None:
 76    binary = _find_chromium()
 77    if not binary:
 78        raise ChromiumError(
 79            "No chromium binary found on PATH (tried: chromium, "
 80            "chromium-browser, google-chrome)"
 81        )
 82    cmd = [binary, *BASE_FLAGS, *args]
 83    try:
 84        subprocess.run(cmd, check=True, capture_output=True, timeout=timeout)
 85    except subprocess.TimeoutExpired as exc:
 86        raise ChromiumError(f"chromium timed out after {timeout}s") from exc
 87    except subprocess.CalledProcessError as exc:
 88        stderr = (exc.stderr or b"").decode("utf-8", errors="replace").strip()
 89        raise ChromiumError(
 90            f"chromium exited {exc.returncode}: {stderr or '(no stderr)'}"
 91        ) from exc
 92
 93
 94def _save(source: Path, filename: str) -> str:
 95    if default_storage.exists(filename):
 96        default_storage.delete(filename)
 97    with source.open("rb") as fp:
 98        default_storage.save(filename, File(fp))
 99    return default_storage.url(filename)
100
101
102def generate_screenshot_from_url(
103    url: str,
104    filename: str,
105    *,
106    viewport: Tuple[int, int] = DEFAULT_VIEWPORT,
107    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
108    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
109) -> str:
110    with _tempfile(".png") as out:
111        _run(
112            [
113                f"--screenshot={out}",
114                f"--window-size={viewport[0]},{viewport[1]}",
115                f"--virtual-time-budget={virtual_time_budget_ms}",
116                url,
117            ],
118            timeout,
119        )
120        return _save(out, filename)
121
122
123def generate_screenshot_from_html(
124    html: str,
125    filename: str,
126    *,
127    viewport: Tuple[int, int] = DEFAULT_VIEWPORT,
128    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
129    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
130) -> str:
131    with _html_tempfile(html) as url:
132        return generate_screenshot_from_url(
133            url,
134            filename,
135            viewport=viewport,
136            virtual_time_budget_ms=virtual_time_budget_ms,
137            timeout=timeout,
138        )
139
140
141def generate_pdf_from_url(
142    url: str,
143    filename: str,
144    *,
145    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
146    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
147) -> str:
148    with _tempfile(".pdf") as out:
149        _run(
150            [
151                f"--print-to-pdf={out}",
152                "--no-pdf-header-footer",
153                f"--virtual-time-budget={virtual_time_budget_ms}",
154                url,
155            ],
156            timeout,
157        )
158        return _save(out, filename)
159
160
161def generate_pdf_from_html(
162    html: str,
163    filename: str,
164    *,
165    virtual_time_budget_ms: int = DEFAULT_VIRTUAL_TIME_BUDGET_MS,
166    timeout: int = DEFAULT_SUBPROCESS_TIMEOUT_S,
167) -> str:
168    with _html_tempfile(html) as url:
169        return generate_pdf_from_url(
170            url,
171            filename,
172            virtual_time_budget_ms=virtual_time_budget_ms,
173            timeout=timeout,
174        )