spectre_pdf 1.0.0

"""Head-to-head: spectre-rs v0.5 positional surface vs pymupdf.

Builds the same synthetic 20-page PDF the criterion bench uses, then
times each spectre-rs API against its pymupdf equivalent in-process.
Median of 5 warm runs per call.

Usage:

    python scripts/bench_positional_vs_pymupdf.py [path/to/your.pdf]

If you pass a path, the bench runs against that file instead of the
synthetic one. The synthetic doc is large enough to surface meaningful
numbers (it's the same fixture criterion's `--quick` mode uses) but a
real document is a better calibration.

Requires `spectre_rs` built with the Python feature
(`maturin develop --release --features python`) and `pymupdf`.
"""
from __future__ import annotations

import statistics
import sys
import time
from pathlib import Path

import spectre_rs  # type: ignore
import fitz  # type: ignore


def build_synthetic_pdf(n_pages: int = 20) -> bytes:
    """Mirror benches/positional.rs::build_synthetic_pdf."""
    doc = fitz.open()
    for i in range(n_pages):
        page = doc.new_page(width=612, height=792)
        for line in range(25):
            y = 740 - line * 14
            page.insert_text(
                (72, 792 - y),
                f"Page {i+1} line {line+1} the quick brown fox jumps over the lazy dog.",
                fontname="helv",
                fontsize=12,
            )
    return doc.tobytes()


def time_fn(fn, *, warmup: int = 2, repeat: int = 5) -> float:
    for _ in range(warmup):
        fn()
    samples = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        fn()
        samples.append(time.perf_counter() - t0)
    return statistics.median(samples)


def spectre_doc_words(doc):
    return doc.words()


def spectre_doc_blocks(doc):
    return doc.blocks()


def spectre_doc_search(doc):
    return doc.search("fox")


def spectre_doc_toc(doc):
    return doc.toc()


def spectre_doc_links(doc):
    return doc.links()


def spectre_doc_annotations(doc):
    return doc.annotations()


def spectre_doc_images(doc):
    return doc.images()


def spectre_doc_pages(doc):
    return doc.pages()


def spectre_doc_info(doc):
    return doc.info()


def spectre_words(pdf_bytes: bytes):
    return spectre_rs.extract_words(pdf_bytes)


def pymupdf_words(pdf_bytes: bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            out.extend(page.get_text("words"))
    return out


def spectre_blocks(pdf_bytes: bytes):
    return spectre_rs.extract_blocks(pdf_bytes)


def pymupdf_blocks(pdf_bytes: bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            out.extend(page.get_text("blocks"))
    return out


def spectre_search(pdf_bytes: bytes):
    return spectre_rs.search(pdf_bytes, "fox")


def pymupdf_search(pdf_bytes: bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            out.extend(page.search_for("fox"))
    return out


def spectre_toc(pdf_bytes: bytes):
    return spectre_rs.extract_toc(pdf_bytes)


def pymupdf_toc(pdf_bytes: bytes):
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        return doc.get_toc()


def spectre_links(pdf_bytes: bytes):
    return spectre_rs.extract_links(pdf_bytes)


def pymupdf_links(pdf_bytes: bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            out.extend(page.get_links())
    return out


def spectre_images(pdf_bytes: bytes):
    return spectre_rs.extract_images(pdf_bytes)


def pymupdf_images(pdf_bytes: bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            out.extend(page.get_images())
    return out


def spectre_page_info(pdf_bytes: bytes):
    return spectre_rs.extract_page_info(pdf_bytes)


def pymupdf_page_info(pdf_bytes: bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            out.append((page.rect, page.rotation, page.mediabox))
    return out


def spectre_doc_info(pdf_bytes: bytes):
    return spectre_rs.extract_document_info(pdf_bytes)


def pymupdf_doc_info(pdf_bytes: bytes):
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        return (doc.page_count, doc.pdf_version(), doc.is_encrypted)


def spectre_annotations(pdf_bytes: bytes):
    return spectre_rs.extract_annotations(pdf_bytes)


def pymupdf_annotations(pdf_bytes: bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page in doc:
            out.extend(list(page.annots() or []))
    return out


def main() -> int:
    if len(sys.argv) > 1:
        path = Path(sys.argv[1])
        pdf_bytes = path.read_bytes()
        label = str(path)
    else:
        pdf_bytes = build_synthetic_pdf(20)
        label = "synthetic 20-page prose PDF"

    print(f"corpus: {label}  ({len(pdf_bytes)} bytes)\n")

    # Free-function comparison (each call re-parses the PDF — apples to
    # apples vs pymupdf opening a fresh Document every call).
    pairs_free = [
        ("words (free)", spectre_words, pymupdf_words),
        ("blocks (free)", spectre_blocks, pymupdf_blocks),
        ("search('fox') (free)", spectre_search, pymupdf_search),
        ("toc (free)", spectre_toc, pymupdf_toc),
        ("links (free)", spectre_links, pymupdf_links),
        ("annotations (free)", spectre_annotations, pymupdf_annotations),
        ("images (free)", spectre_images, pymupdf_images),
        ("page_info (free)", spectre_page_info, pymupdf_page_info),
        ("doc_info (free)", spectre_doc_info, pymupdf_doc_info),
    ]

    # Persistent-handle comparison: both libraries open the PDF once and
    # the timed call only does the extraction. This is the apples-to-apples
    # shape — pymupdf's fitz.Document is exactly this pattern.
    s_doc = spectre_rs.Document(pdf_bytes)
    p_doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    pairs_doc = [
        ("words (handle)", lambda: spectre_doc_words(s_doc), lambda: [w for page in p_doc for w in page.get_text("words")]),
        ("blocks (handle)", lambda: spectre_doc_blocks(s_doc), lambda: [b for page in p_doc for b in page.get_text("blocks")]),
        ("search('fox') (handle)", lambda: spectre_doc_search(s_doc), lambda: [r for page in p_doc for r in page.search_for("fox")]),
        ("toc (handle)", lambda: spectre_doc_toc(s_doc), lambda: p_doc.get_toc()),
        ("links (handle)", lambda: spectre_doc_links(s_doc), lambda: [l for page in p_doc for l in page.get_links()]),
        ("annotations (handle)", lambda: spectre_doc_annotations(s_doc), lambda: [a for page in p_doc for a in (page.annots() or [])]),
        ("images (handle)", lambda: spectre_doc_images(s_doc), lambda: [i for page in p_doc for i in page.get_images()]),
        ("pages (handle)", lambda: spectre_doc_pages(s_doc), lambda: [(page.rect, page.rotation, page.mediabox) for page in p_doc]),
        ("info (handle)", lambda: s_doc.info(), lambda: (p_doc.page_count, p_doc.pdf_version(), p_doc.is_encrypted)),
    ]

    print(f"{'api':<28}{'spectre (ms)':>14}{'pymupdf (ms)':>16}{'speedup':>14}")
    print("-" * 72)
    for batch_label, batch in [("--- free-function pattern (re-parse each call) ---", pairs_free),
                                ("--- persistent-handle pattern (parse once) ---", pairs_doc)]:
        print(batch_label)
        for name, sfn, pfn in batch:
            try:
                s = time_fn(sfn if isinstance(sfn, type(lambda: 0)) and sfn.__code__.co_argcount == 0 else lambda: sfn(pdf_bytes)) * 1000
            except Exception as e:  # noqa: BLE001
                print(f"{name:<28}{'spectre err: ' + str(e)[:40]:>40}")
                continue
            try:
                p = time_fn(pfn if isinstance(pfn, type(lambda: 0)) and pfn.__code__.co_argcount == 0 else lambda: pfn(pdf_bytes)) * 1000
            except Exception as e:  # noqa: BLE001
                print(f"{name:<28}{s:>14.2f}{'pymupdf err':>16}{'-':>14}")
                continue
            speedup = p / s if s > 0 else float("inf")
            print(f"{name:<28}{s:>14.2f}{p:>16.2f}{speedup:>13.2f}x")
    return 0


if __name__ == "__main__":
    sys.exit(main())