cs249r_book/tools/audit/audit_math_rendering.py

#!/usr/bin/env python3
"""
Math-rendering audit for MLSysBook HTML output.

For each chapter, builds the HTML (via `./book/binder build html <chap>`),
then scans the rendered HTML for raw LaTeX leakage that escaped MathJax.

A "leak" is a LaTeX command or pattern that appears in user-visible prose
(outside <script>, <style>, <code>, <pre>, or <span class="math ...">
elements). These render as literal text in browsers and indicate a bug.

Usage (run from repo root):
    python3 tools/audit/audit_math_rendering.py               # all chapters
    python3 tools/audit/audit_math_rendering.py vol1/intro    # one chapter
    python3 tools/audit/audit_math_rendering.py --skip-build  # just scan
"""

from __future__ import annotations
import argparse
import json
import re
import subprocess
import sys
import time
from dataclasses import dataclass, field, asdict
from html.parser import HTMLParser
from pathlib import Path
from typing import Iterable

REPO = Path(__file__).resolve().parents[2]
BUILD_DIRS = {
    "vol1": REPO / "book" / "quarto" / "_build" / "html-vol1",
    "vol2": REPO / "book" / "quarto" / "_build" / "html-vol2",
}
BINDER = REPO / "book" / "binder"

# Patterns that, if found in user-visible HTML text, indicate a rendering bug.
# Order matters for reporting; most diagnostic first.
LEAK_PATTERNS: list[tuple[str, re.Pattern]] = [
    ("backslash-cmd", re.compile(r"\\(?:times|frac|approx|alpha|beta|gamma|delta|sigma|mu|eta|lambda|theta|sqrt|sum|int|cdot|leq|geq|neq|partial|nabla|infty|log|ln|exp|min|max|pm|mp|forall|exists|in|notin|to|gets|mapsto|rightarrow|Rightarrow|prod|approx|equiv|sim|propto|land|lor|neg|text|mathbf|mathit|mathcal|mathbb)\b")),  # codespell:ignore notin
    ("caret-brace", re.compile(r"(?<![A-Za-z0-9_])\^\{[^}\n]{1,40}\}")),
    ("under-brace", re.compile(r"(?<![A-Za-z0-9_])_\{[^}\n]{1,40}\}")),
    ("bare-pow10", re.compile(r"\b10\^\{[+-]?\d{1,4}\}")),
    ("dollar-math-text", re.compile(r"\$\\[a-zA-Z]+|\\[a-zA-Z]+\$")),
]

SKIP_TEXT_TAGS = {"script", "style", "pre", "code", "mjx-container", "math", "annotation", "head"}


class VisibleTextExtractor(HTMLParser):
    """Collect user-visible text while skipping HTML/MathJax code zones."""

    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self._skip_stack: list[str] = []
        self._chunks: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        tag = tag.lower()
        attr_map = {name.lower(): value or "" for name, value in attrs}
        classes = set(attr_map.get("class", "").split())
        if tag in SKIP_TEXT_TAGS or (tag == "span" and "math" in classes):
            self._skip_stack.append(tag)

    def handle_endtag(self, tag: str) -> None:
        tag = tag.lower()
        if self._skip_stack and self._skip_stack[-1] == tag:
            self._skip_stack.pop()

    def handle_data(self, data: str) -> None:
        if not self._skip_stack:
            self._chunks.append(data)

    def text(self) -> str:
        return " ".join(self._chunks)


def visible_text_from_html(html: str) -> str:
    parser = VisibleTextExtractor()
    parser.feed(html)
    parser.close()
    return parser.text()


@dataclass
class Leak:
    pattern: str
    match: str
    context: str
    char_offset: int


@dataclass
class ChapterReport:
    name: str
    volume: str
    qmd_path: str
    build_ok: bool
    build_seconds: float
    html_path: str | None = None
    leaks: list[Leak] = field(default_factory=list)
    error: str | None = None


def list_chapters() -> list[tuple[str, str, Path]]:
    """Return list of (name, volume, qmd_path) for all chapters in vol1+vol2."""
    chapters = []
    for vol in ("vol1", "vol2"):
        contents = REPO / "book" / "quarto" / "contents" / vol
        for qmd in sorted(contents.glob("*/*.qmd")):
            # Skip non-chapter files (parts, frontmatter sub-files starting with _)
            if qmd.name.startswith("_") and qmd.name not in (
                "_notation_body.qmd",
                "_notation_distributed.qmd",
            ):
                continue
            # Use the chapter slug (filename stem) -- binder accepts that
            chapters.append((qmd.stem, vol, qmd))
    return chapters


def build_chapter(name: str, volume: str) -> tuple[bool, float, str]:
    """Run `./book/binder build html <vol>/<name>` and return (ok, secs, output)."""
    chap = f"{volume}/{name}"
    t0 = time.time()
    proc = subprocess.run(
        [str(BINDER), "build", "html", chap],
        cwd=REPO,
        capture_output=True,
        text=True,
        timeout=600,
    )
    secs = time.time() - t0
    ok = proc.returncode == 0 and "build completed" in (proc.stdout or "").lower()
    return ok, secs, (proc.stdout or "") + (proc.stderr or "")


def find_html(name: str, volume: str, qmd_path: Path) -> Path | None:
    """Locate the rendered HTML for this chapter."""
    base = BUILD_DIRS[volume]
    rel = qmd_path.relative_to(REPO / "book" / "quarto").with_suffix(".html")
    candidate = base / rel
    if candidate.exists():
        return candidate
    # Fallback: search by filename stem
    for hit in base.rglob(f"{name}.html"):
        return hit
    return None


def scan_html(html: str) -> list[Leak]:
    text = visible_text_from_html(html)

    leaks: list[Leak] = []
    for label, pat in LEAK_PATTERNS:
        for m in pat.finditer(text):
            start = max(0, m.start() - 40)
            end = min(len(text), m.end() + 40)
            ctx = text[start:end].replace("\n", " ").strip()
            leaks.append(Leak(label, m.group(0), ctx, m.start()))
            if len(leaks) > 200:  # cap per-chapter
                return leaks
    return leaks


def audit_chapter(name: str, volume: str, qmd_path: Path, skip_build: bool) -> ChapterReport:
    rep = ChapterReport(
        name=name,
        volume=volume,
        qmd_path=str(qmd_path.relative_to(REPO)),
        build_ok=True,
        build_seconds=0.0,
    )
    if not skip_build:
        try:
            ok, secs, output = build_chapter(name, volume)
        except subprocess.TimeoutExpired:
            rep.build_ok = False
            rep.error = "build timed out"
            return rep
        rep.build_ok = ok
        rep.build_seconds = secs
        if not ok:
            rep.error = (output[-2000:] if output else "no output")
            return rep

    html_path = find_html(name, volume, qmd_path)
    if html_path is None:
        rep.error = "html output not found"
        return rep
    rep.html_path = str(html_path.relative_to(REPO))
    try:
        html = html_path.read_text(encoding="utf-8", errors="replace")
    except Exception as e:
        rep.error = f"read error: {e}"
        return rep
    rep.leaks = scan_html(html)
    return rep


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("chapters", nargs="*", help="vol1/<chap> tokens to audit; default = all")
    ap.add_argument("--skip-build", action="store_true", help="reuse existing _build HTML")
    ap.add_argument("--report", default="audit-math-report.json")
    ap.add_argument("--text-report", default="audit-math-report.md")
    args = ap.parse_args()

    all_chapters = list_chapters()
    if args.chapters:
        wanted = set(args.chapters)
        chapters = [c for c in all_chapters if f"{c[1]}/{c[0]}" in wanted or c[0] in wanted]
    else:
        chapters = all_chapters

    print(f"Auditing {len(chapters)} chapters", flush=True)
    reports: list[ChapterReport] = []
    for i, (name, vol, qmd) in enumerate(chapters, 1):
        print(f"[{i:3d}/{len(chapters)}] {vol}/{name} ... ", end="", flush=True)
        try:
            rep = audit_chapter(name, vol, qmd, args.skip_build)
        except Exception as e:
            print(f"EXCEPTION: {e}", flush=True)
            rep = ChapterReport(
                name=name, volume=vol, qmd_path=str(qmd.relative_to(REPO)),
                build_ok=False, build_seconds=0.0, error=f"exception: {e}",
            )
        if not rep.build_ok:
            print(f"BUILD FAILED ({rep.error[:80] if rep.error else '?'})", flush=True)
        elif rep.error:
            print(f"ERROR: {rep.error[:80]}", flush=True)
        elif rep.leaks:
            counts: dict[str, int] = {}
            for L in rep.leaks:
                counts[L.pattern] = counts.get(L.pattern, 0) + 1
            print(f"{len(rep.leaks)} leak(s): {dict(counts)} ({rep.build_seconds:.0f}s)", flush=True)
        else:
            print(f"clean ({rep.build_seconds:.0f}s)", flush=True)
        reports.append(rep)

    # JSON dump
    Path(args.report).write_text(json.dumps([asdict(r) for r in reports], indent=2))

    # Human-readable report
    lines = ["# Math Rendering Audit Report", "",
             f"Total chapters audited: **{len(reports)}**", ""]
    failed = [r for r in reports if not r.build_ok]
    leaky = [r for r in reports if r.build_ok and r.leaks]
    clean = [r for r in reports if r.build_ok and not r.leaks and not r.error]
    errored = [r for r in reports if r.build_ok and r.error]

    lines += [f"- Clean: **{len(clean)}**",
              f"- Leaky: **{len(leaky)}**",
              f"- Build failures: **{len(failed)}**",
              f"- Other errors: **{len(errored)}**", ""]

    if leaky:
        lines += ["## Chapters with LaTeX leakage", ""]
        for r in leaky:
            counts: dict[str, int] = {}
            for L in r.leaks:
                counts[L.pattern] = counts.get(L.pattern, 0) + 1
            lines += [f"### `{r.volume}/{r.name}` ({len(r.leaks)} leaks: {dict(counts)})",
                      f"- qmd: `{r.qmd_path}`",
                      f"- html: `{r.html_path}`",
                      ""]
            shown: dict[str, int] = {}
            for L in r.leaks[:30]:
                if shown.get(L.pattern, 0) >= 5:
                    continue
                shown[L.pattern] = shown.get(L.pattern, 0) + 1
                lines += [f"  - **{L.pattern}** match `{L.match}`",
                          f"    > ...{L.context}..."]
            lines.append("")

    if failed:
        lines += ["## Build failures", ""]
        for r in failed:
            lines += [f"- `{r.volume}/{r.name}`: {(r.error or '?')[:200]}"]
        lines.append("")

    if errored:
        lines += ["## Other errors", ""]
        for r in errored:
            lines += [f"- `{r.volume}/{r.name}`: {(r.error or '?')[:200]}"]
        lines.append("")

    if clean:
        lines += ["## Clean chapters", ""]
        for r in clean:
            lines += [f"- `{r.volume}/{r.name}` ({r.build_seconds:.0f}s)"]
        lines.append("")

    Path(args.text_report).write_text("\n".join(lines))
    print(f"\nReports written: {args.report}, {args.text_report}", flush=True)
    print(f"Summary: {len(clean)} clean | {len(leaky)} leaky | "
          f"{len(failed)} build-failed | {len(errored)} errored", flush=True)
    sys.exit(0 if not leaky and not failed else 1)


if __name__ == "__main__":
    main()