pdf2md/tests/conftest.py

"""Synthetische Test-PDFs: digital, gescannt, gescannt mit OCR-Ebene, gemischt."""

import random
import shutil
import sys
from pathlib import Path

import pymupdf
import pytest

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

TITLE = "Das Testbuch"

CSS = """
body { font-family: serif; font-size: 12pt; }
h1 { font-size: 24pt; font-weight: bold; margin-bottom: 12pt; }
h2 { font-size: 18pt; font-weight: bold; margin-top: 14pt; margin-bottom: 10pt; }
p { margin-top: 0; margin-bottom: 8pt; text-align: left; }
"""

PAGES = [
    f"""
    <h1>{TITLE}</h1>
    <p>Ein kurzer Roman zum Testen der Umwandlung.</p>
    <h2>Inhaltsverzeichnis</h2>
    <p>Teil I: Der Anfang<br>1 Der Notar<br>2 Ein Brief</p>
    """,
    """
    <h2>Teil I: Der Anfang</h2>
    <h2>1 Der Notar</h2>
    <p>Albrecht van Oie stand am Fenster und blickte auf die nasse Straße
    hinunter. Der Regen hatte am Morgen begonnen und wollte nicht aufhören.</p>
    <p>Er dachte an die alten Zeiten in Berlin, an die Freunde von damals und
    an die vielen Briefe, die er nie beantwortet hatte.</p>
    <p>Am Nachmittag ging er zum Notar in der Kantstraße. Dort sprach er
    lange mit dem Geheim-</p>
    """,
    """
    <p>dienst über die Lage in der Stadt und über die Zukunft des Landes.</p>
    <h2>2 Ein Brief</h2>
    <p>Der Brief lag auf dem Tisch. Oie nahm ihn langsam zur Hand und las die
    ersten Zeilen mit wachsender Spannung.</p>
    <p>»Teurer Albrecht!« stand da in einer schönen alten Handschrift.</p>
    """,
    """
    <p>Am Abend saß er lange in der Küche und trank einen Tee. Die Gedanken
    an den Brief ließen ihn nicht los.</p>
    <p>Schließlich ging er zu Bett und schlief tief und fest bis zum
    nächsten Morgen.</p>
    """,
]

OUTLINE = [
    [1, TITLE, 1],
    [2, "Teil I: Der Anfang", 2],
    [2, "1 Der Notar", 2],
    [2, "2 Ein Brief", 3],
]

needs_tesseract = pytest.mark.skipif(
    not (pymupdf.get_tessdata() and (Path(pymupdf.get_tessdata()) / "deu.traineddata").is_file()),
    reason="Tesseract mit deutschen Sprachdaten nicht installiert",
)
needs_pandoc = pytest.mark.skipif(
    shutil.which("pandoc") is None, reason="pandoc nicht installiert"
)


def _make_digital() -> pymupdf.Document:
    doc = pymupdf.open()
    for number, html in enumerate(PAGES, start=1):
        page = doc.new_page(width=420, height=595)
        body = pymupdf.Rect(50, 60, 370, 540)
        spare, _scale = page.insert_htmlbox(body, html, css=CSS, scale_low=1)
        assert spare >= 0, f"Testseite {number} zu lang"
        if number > 1:
            # Kolumnentitel und Seitenzahl wie in einem gedruckten Buch
            page.insert_text((170, 35), TITLE, fontsize=9)
            page.insert_text((205, 570), str(number), fontsize=9)
    # Ein Bild auf Seite 3
    pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 120, 80), 0)
    pix.set_rect(pix.irect, (40, 90, 160))
    pix.set_rect(pymupdf.IRect(20, 20, 100, 60), (230, 200, 60))
    doc[2].insert_image(pymupdf.Rect(150, 420, 270, 500), pixmap=pix)
    doc.set_toc(OUTLINE)
    doc.set_metadata({"title": TITLE, "author": "Erika Muster"})
    return doc


def _scan_page(target: pymupdf.Document, page: pymupdf.Page) -> None:
    pix = page.get_pixmap(dpi=300, colorspace=pymupdf.csGRAY)
    new = target.new_page(width=page.rect.width, height=page.rect.height)
    new.insert_image(new.rect, pixmap=pix)


def _ocr_page(target: pymupdf.Document, page: pymupdf.Page) -> None:
    # Achtung: MuPDFs OCR liefert bei Graustufen-Pixmaps keinen Text -> RGB.
    pix = page.get_pixmap(dpi=300)
    with pymupdf.open("pdf", pix.pdfocr_tobytes(language="deu")) as ocr:
        target.insert_pdf(ocr)


@pytest.fixture(scope="session")
def pdfs(tmp_path_factory) -> dict[str, Path]:
    folder = tmp_path_factory.mktemp("pdfs")
    paths = {}

    digital = _make_digital()
    paths["digital"] = folder / "digital.pdf"
    digital.save(paths["digital"])

    scanned = pymupdf.open()
    for page in digital:
        _scan_page(scanned, page)
    paths["scanned"] = folder / "scanned.pdf"
    scanned.save(paths["scanned"])

    mixed = pymupdf.open()
    mixed.insert_pdf(digital, from_page=0, to_page=1)
    for page in list(digital)[2:]:
        _scan_page(mixed, page)
    paths["mixed"] = folder / "mixed.pdf"
    mixed.save(paths["mixed"])

    # Digitales PDF mit einer reinen Fotoseite (Rauschen) am Ende
    photo = pymupdf.open()
    photo.insert_pdf(digital)
    rng = random.Random(7)
    noise = pymupdf.Pixmap(
        pymupdf.csRGB, 200, 280, bytes(rng.randrange(256) for _ in range(200 * 280 * 3)), 0
    )
    page = photo.new_page(width=420, height=595)
    page.insert_image(page.rect, pixmap=noise)
    paths["photo"] = folder / "photo.pdf"
    photo.save(paths["photo"])

    blank = pymupdf.open()
    blank.new_page()
    paths["blank"] = folder / "blank.pdf"
    blank.save(paths["blank"])

    # Scan mit falscher (unsichtbarer) OCR-Ebene, wie von einem schlechten OCR-Lauf
    bad_ocr = pymupdf.open()
    for page in digital:
        _scan_page(bad_ocr, page)
        for y in range(80, 500, 20):
            bad_ocr[-1].insert_text((60, y), "Falscher Text " * 4, fontsize=10, render_mode=3)
    paths["bad_ocr"] = folder / "bad_ocr.pdf"
    bad_ocr.save(paths["bad_ocr"])

    if pymupdf.get_tessdata():
        scanned_ocr = pymupdf.open()
        for page in digital:
            _ocr_page(scanned_ocr, page)
        paths["scanned_ocr"] = folder / "scanned_ocr.pdf"
        scanned_ocr.save(paths["scanned_ocr"])

    return paths