pdf2md/tests/test_split_pages.py

"""--split seiten, --page-ranges, Nachspann, --split-chars, --from-md --split, --epub-per-file."""

import zipfile

import pytest

import pdf2md as p
from conftest import needs_pandoc, needs_tesseract
from test_e2e import run
from test_split import BOOK

# Rohtext je Seite (wie von pymupdf4llm) und das daraus normalisierte Markdown:
# Absatz B läuft von Seite 2 auf Seite 3 (mit Silbentrennung), auf Seite 3
# steht vor Absatz D eine Überschrift.
PAGE_TEXTS = {
    0: "Absatz A steht ganz auf der ersten Seite.\n",
    1: "Absatz B beginnt auf der zweiten Seite und reicht mit einem Geheim-\n",
    2: "dienst bis auf die dritte Seite.\n\n## Neues Kapitel\n\nAbsatz D folgt danach.\n",
    3: "Absatz E auf der vierten Seite.\n",
}
MARKDOWN = (
    "Absatz A steht ganz auf der ersten Seite.\n\n"
    "Absatz B beginnt auf der zweiten Seite und reicht mit einem Geheimdienst bis auf "
    "die dritte Seite.\n\n"
    "## Neues Kapitel\n\n"
    "Absatz D folgt danach.\n\n"
    "Absatz E auf der vierten Seite.\n"
)


def _flat(text: str) -> str:
    return text.replace("\n", "")


def test_md_blocks_keep_text():
    assert "".join(p.md_blocks(BOOK)) == BOOK
    assert "".join(p.md_blocks(MARKDOWN)) == MARKDOWN


def test_block_pages_start_and_end():
    spans = p.block_pages(p.md_blocks(MARKDOWN), PAGE_TEXTS)
    assert spans == [(0, 0), (1, 2), (2, 2), (2, 2), (3, 3)]


def test_block_pages_inherit_unknown_blocks():
    blocks = ["Absatz A steht ganz auf der ersten Seite.\n\n", "![](images/x.jpg)\n\n", "Ja.\n\n",
              "Absatz E auf der vierten Seite.\n"]
    assert [s for s, _ in p.block_pages(blocks, PAGE_TEXTS)] == [0, 0, 0, 3]


def test_split_by_pages_keeps_paragraph_over_page_break():
    sections = p.split_by_pages(MARKDOWN, [0, 2], PAGE_TEXTS, align=False)
    # Absatz B beginnt auf Seite 2 -> bleibt vorn; geschnitten wird vor der
    # Überschrift, die auf Seite 3 beginnt.
    assert sections[0]["body"].endswith("die dritte Seite.\n")
    assert sections[0]["source_pages"] == [1, 3]
    assert sections[1]["body"].startswith("## Neues Kapitel")
    assert sections[1]["title"] == "Neues Kapitel" and sections[1]["source_pages"] == [3, 4]
    assert _flat("".join(s["body"] for s in sections)) == _flat(MARKDOWN)


def test_split_by_pages_heading_moves_to_next_file():
    # Grenze vor Seite 3: der erste dort beginnende Block ist Absatz D; die
    # Überschrift davor gehört zur neuen Datei.
    texts = PAGE_TEXTS | {2: "dienst bis auf die dritte Seite.\n\n## Neues Kapitel\n",
                          3: "Absatz D folgt danach.\n\nAbsatz E auf der vierten Seite.\n"}
    sections = p.split_by_pages(MARKDOWN, [0, 3], texts, align=False)
    assert sections[1]["body"].startswith("## Neues Kapitel\n\nAbsatz D")
    assert not sections[0]["body"].rstrip().endswith("Kapitel")


def test_split_by_pages_aligns_to_nearby_heading():
    # Grenze vor Seite 4 (Absatz E); die Überschrift auf Seite 3 liegt eine Seite davor.
    aligned = p.split_by_pages(MARKDOWN, [0, 3], PAGE_TEXTS, align=True)
    strict = p.split_by_pages(MARKDOWN, [0, 3], PAGE_TEXTS, align=False)
    assert aligned[1]["body"].startswith("## Neues Kapitel")
    assert strict[1]["body"].startswith("Absatz E") and strict[1]["title"] == "Seite 4"


@pytest.mark.parametrize("starts", [[0, 1], [0, 1, 2, 3], [0, 2], [0]])
def test_split_by_pages_loses_nothing(starts):
    for align in (True, False):
        sections = p.split_by_pages(MARKDOWN, starts, PAGE_TEXTS, align=align)
        assert _flat("".join(s["body"] for s in sections)) == _flat(MARKDOWN)


def test_parse_page_ranges():
    assert p.parse_page_ranges("1-12,13-40,41-", 100) == [0, 12, 40]
    assert p.parse_page_ranges("5-10, 20-", 30) == [0, 4, 10, 19]  # Lücken = eigene Dateien
    assert p.parse_page_ranges("3", 10) == [0, 2, 3]
    for bad in ("1-12,10-20", "0-5", "5-200", "a-b", "8-3"):
        with pytest.raises(ValueError):
            p.parse_page_ranges(bad, 100)


def test_backmatter_becomes_own_file():
    book = BOOK.replace("## Impressum\n\nAlle Rechte vorbehalten.\n",
                        "### Der Autor\n\nLebt in Berlin.\n\n### Impressum\n\nRechte.\n")
    sections = p.split_markdown(book, "teile")
    assert [s["title"] for s in sections] == ["Vorspann", "Teil I: Der Anfang", "Nachspann"]
    assert sections[-1]["body"].startswith("### Der Autor") and "Impressum" in sections[-1]["body"]
    assert _flat("".join(s["body"] for s in sections)) == _flat(book)


def test_split_chars_and_heading_titles():
    paragraphs = [f"Absatz {n} " + " ".join(["Text"] * 40) for n in range(20)]
    markdown = "## Erster Teil\n\n" + "\n\n".join(paragraphs) + "\n"
    sections = p.split_markdown(markdown, "kapitel", size=1500)
    assert sections[0]["title"] == "Erster Teil" and sections[1]["title"] == "Abschnitt 2"
    assert len(p.split_markdown(markdown, "kapitel", size=100_000)) == 1


def test_rebase_image_links(tmp_path):
    md = "![](buch_images/a.jpg) ![x](https://e.de/b.png) ![](/abs/c.jpg)"
    out = p.rebase_image_links(md, tmp_path, tmp_path / "buch")
    assert out == "![](../buch_images/a.jpg) ![x](https://e.de/b.png) ![](/abs/c.jpg)"


# --- Ende zu Ende ------------------------------------------------------------------

def _heads(folder):
    return [p.split_frontmatter(f.read_text())[0] for f in sorted(folder.glob("0*.md"))]


def test_cli_split_pages(pdfs, tmp_path):
    run(pdfs["digital"], tmp_path / "buch.md", "--split", "seiten", "--pages-per-file", "1")
    heads = _heads(tmp_path / "buch")
    assert [h["source_pages"] for h in heads] == [[1, 1], [2, 3], [3, 3], [4, 4]]
    assert [h["title"] for h in heads] == ["Das Testbuch", "Teil I: Der Anfang", "2 Ein Brief", "Seite 4"]
    whole = tmp_path / "ganz.md"
    run(pdfs["digital"], whole)
    body = "\n\n".join(p.split_frontmatter(f.read_text())[1].strip()
                       for f in sorted((tmp_path / "buch").glob("0*.md")))
    assert _flat(body) == _flat(p.split_frontmatter(whole.read_text())[1].strip())


def test_cli_page_ranges(pdfs, tmp_path):
    run(pdfs["digital"], tmp_path / "buch.md", "--split", "seiten", "--page-ranges", "2-3")
    assert [h["source_pages"][0] for h in _heads(tmp_path / "buch")] == [1, 2, 4]


@needs_tesseract
def test_cli_split_pages_scanned(pdfs, tmp_path):
    # Seitenzuordnung über den OCR-Text (die Scanseiten haben selbst keinen Text)
    run(pdfs["scanned"], tmp_path / "buch.md", "--split", "seiten", "--pages-per-file", "2",
        "--jobs", "1")
    heads = _heads(tmp_path / "buch")
    assert len(heads) == 2 and heads[1]["source_pages"][0] == 3


def test_cli_split_options_are_checked(pdfs, tmp_path):
    import subprocess
    import sys
    for extra in (["--pages-per-file", "3"], ["--split", "kapitel", "--page-ranges", "1-2"],
                  ["--epub-per-file"]):
        done = subprocess.run([sys.executable, str(p.__file__), str(pdfs["digital"]),
                               str(tmp_path / "x.md"), *extra], capture_output=True, text=True)
        assert done.returncode == 2 and "nur mit --split" in done.stderr


@needs_pandoc
def test_from_md_split_and_epub_per_file(pdfs, tmp_path):
    md = tmp_path / "buch.md"
    run(pdfs["digital"], md, "--images")
    run("--from-md", md, "--split", "kapitel", "--epub-per-file")
    folder = tmp_path / "buch"
    names = sorted(f.name for f in folder.glob("*.md"))
    assert "index.md" in names and any(n.endswith("_1-der-notar.md") for n in names)
    text = "".join(f.read_text() for f in folder.glob("0*.md"))
    assert "](../buch_images/" in text
    meta, _ = p.split_frontmatter((folder / "index.md").read_text())
    assert meta["cover"] == "../buch_images/cover.jpg"
    epubs = sorted((folder / "epub").glob("*.epub"))
    assert len(epubs) == len(names) - 1
    notar = next(e for e in epubs if "1-der-notar" in e.name)
    with zipfile.ZipFile(notar) as z:
        opf = z.read("EPUB/content.opf").decode()
        images = [n for n in z.namelist() if n.endswith(".jpg")]
    assert "Das Testbuch - 1 Der Notar" in opf and "cover" in opf and images
    assert not (tmp_path / "buch.epub").exists()  # Gesamt-EPUB nur mit --epub