pdf2md/tests/test_classify.py

"""Seitenklassifizierung und OCR-Planung."""

import pymupdf
import pytest

import pdf2md as p
from conftest import needs_tesseract


def kinds_of(path):
    with pymupdf.open(path) as doc:
        return [p.classify_page(page) for page in doc]


def test_digital_pages_are_text(pdfs):
    assert kinds_of(pdfs["digital"]) == [p.PAGE_TEXT] * 4


def test_scanned_pages_are_scans(pdfs):
    assert kinds_of(pdfs["scanned"]) == [p.PAGE_SCAN] * 4


@needs_tesseract
def test_scans_with_ocr_layer_are_detected(pdfs):
    assert kinds_of(pdfs["scanned_ocr"]) == [p.PAGE_OCR_LAYER] * 4


def test_mixed_document(pdfs):
    assert kinds_of(pdfs["mixed"]) == [p.PAGE_TEXT, p.PAGE_TEXT, p.PAGE_SCAN, p.PAGE_SCAN]


def test_blank_page_is_empty(pdfs):
    assert kinds_of(pdfs["blank"]) == [p.PAGE_EMPTY]


def test_plan_never_uses_no_ocr(pdfs):
    with pymupdf.open(pdfs["scanned"]) as doc:
        _, plan = p.plan_pages(doc, "never", "deu", 300)
    assert plan == [p.RUN_PLAIN] * 4


@needs_tesseract
def test_plan_auto_scanned(pdfs):
    with pymupdf.open(pdfs["scanned"]) as doc:
        _, plan = p.plan_pages(doc, "auto", "deu", 300)
    assert plan == [p.RUN_OCR] * 4


@needs_tesseract
def test_plan_auto_reuses_existing_ocr_layer(pdfs):
    with pymupdf.open(pdfs["scanned_ocr"]) as doc:
        _, plan = p.plan_pages(doc, "auto", "deu", 300)
    assert plan == [p.RUN_REUSE] * 4


def _scan_with_text_layer(path, angle):
    """Scanseite mit unsichtbarer, um `angle` Grad gedrehter Textebene."""
    doc = pymupdf.open()
    page = doc.new_page(width=420, height=595)
    pix = pymupdf.Pixmap(pymupdf.csGRAY, pymupdf.IRect(0, 0, 200, 280), 0)
    pix.set_rect(pix.irect, (230,))
    page.insert_image(page.rect, pixmap=pix)
    for i in range(20):
        point = pymupdf.Point(50, 80 + 20 * i)
        page.insert_text(
            point,
            "Eine Zeile der vorhandenen OCR-Ebene",
            render_mode=3,
            morph=(point, pymupdf.Matrix(1, 1).prerotate(-angle)),
        )
    doc.save(path)
    return path


@needs_tesseract
@pytest.mark.parametrize("angle, expected", [(0, p.RUN_REUSE), (0.5, p.RUN_OCR)])
def test_plan_auto_redoes_skewed_ocr_layer(tmp_path, angle, expected):
    path = _scan_with_text_layer(tmp_path / f"layer_{angle}.pdf", angle)
    with pymupdf.open(path) as doc:
        kinds, plan = p.plan_pages(doc, "auto", "deu", 300)
    assert kinds == [p.PAGE_OCR_LAYER]
    assert plan == [expected]


@needs_tesseract
def test_plan_auto_mixed_probes_scan_pages(pdfs):
    # 2 Text- zu 2 Scanseiten gilt nicht als Scan-Dokument; die Scanseiten
    # werden per Probe-OCR geprüft und bestehen sie.
    with pymupdf.open(pdfs["mixed"]) as doc:
        _, plan = p.plan_pages(doc, "auto", "deu", 300)
    assert plan == [p.RUN_PLAIN, p.RUN_PLAIN, p.RUN_OCR, p.RUN_OCR]


@needs_tesseract
def test_plan_auto_skips_photo_page(pdfs):
    with pymupdf.open(pdfs["photo"]) as doc:
        kinds, plan = p.plan_pages(doc, "auto", "deu", 300)
    assert kinds[-1] == p.PAGE_SCAN
    assert plan == [p.RUN_PLAIN] * 5


@needs_tesseract
def test_plan_always(pdfs):
    with pymupdf.open(pdfs["digital"]) as doc:
        _, plan = p.plan_pages(doc, "always", "deu", 300)
    assert plan == [p.RUN_OCR] * 4


def test_missing_language_gives_clear_error():
    if not pymupdf.get_tessdata():
        pytest.skip("Tesseract nicht installiert")
    with pytest.raises(RuntimeError, match="tesseract-ocr-xyz"):
        p.check_ocr_language("deu+xyz")


@pytest.mark.parametrize(
    "text, language, expected",
    [
        ("Der Brief lag auf dem Tisch. " * 10, "deu", True),
        ("Am Abend saß er lange in der Küche und trank einen Tee. Die Gedanken "
         "an den Brief ließen ihn nicht los.", "deu", True),
        ("The letter was lying on the table in the kitchen, and he read it "
         "twice before he put it away.", "deu+eng", True),
        # Cover mit Titel und Autor: zu wenig Wörter
        ("Reinhard Otto Kranz OPERATION LJUTSCH Band I Schlüssel zur Deutschen Einheit", "deu", False),
        ("fi@ |l1 ;; ¦¦ ~~ ‚, ., " * 20, "deu", False),
        # echte Tesseract-Ausgabe einer Rauschseite
        ("©Se te\nansa |\naa\n.\nerNee =\nae\neea\nBa:Ben\nBena\needaaf\neeee\naar.\nEl eeBR\neen\n", "deu", False),
        ("Zu kurz.", "deu", False),
    ],
)
def test_looks_like_text(text, language, expected):
    assert p.looks_like_text(text, language) is expected


def test_make_tasks_groups_and_chunks():
    plan = [p.RUN_PLAIN] * 3 + [p.RUN_OCR] * 12 + [p.RUN_PLAIN]
    tasks = p.make_tasks(plan)
    assert [(run, pages[0], pages[-1]) for run, pages in tasks] == [
        (p.RUN_PLAIN, 0, 2),
        (p.RUN_OCR, 3, 12),
        (p.RUN_OCR, 13, 14),
        (p.RUN_PLAIN, 15, 15),
    ]