old-books/_pipeline/structure.py

"""Stage 3, structure: find chapter boundaries from the text layer. Audible style file list is made by the script stage."""
import re

import bookio
import pageparts
import pdftext

WORDS = "ONE TWO THREE FOUR FIVE SIX SEVEN EIGHT NINE TEN ELEVEN TWELVE THIRTEEN FOURTEEN FIFTEEN SIXTEEN SEVENTEEN EIGHTEEN NINETEEN TWENTY".split()
ROMAN = {"I": 1, "V": 5, "X": 10, "L": 50}
HEADING_RE = re.compile(r"^\s*CHAPTER\s+([A-Z]+|\d+)\s*$")


def chapter_number(token):
    if token.isdigit():
        return int(token)
    if token in WORDS:
        return WORDS.index(token) + 1
    if set(token) <= set(ROMAN):
        total = 0
        for a, b in zip(token, token[1:] + " "):
            total += -ROMAN[a] if b in ROMAN and ROMAN[b] > ROMAN[a] else ROMAN[a]
        return total
    return None


def find_headings(units):
    """'CHAPTER ONE' then the title line(s) up to a blank line, within the first 12 lines of a unit."""
    found = []
    for ui, u in enumerate(units):
        lines = u["text"].split("\n")
        shown = 0
        for li, ln in enumerate(lines):
            if ln.strip():
                shown += 1
            if shown > 12:
                break
            m = HEADING_RE.match(ln)
            n = chapter_number(m.group(1)) if m else None
            if n is None:
                continue
            k = li + 1
            while k < len(lines) and not lines[k].strip():
                k += 1
            title = []
            while k < len(lines) and lines[k].strip():
                title.append(lines[k].strip())
                k += 1
            found.append({"unit": ui, "line": li, "title_end": k, "number": n,
                          "title": bookio.title_case(" ".join(title))})
            break
    return found


def toc_chapters(units, before_unit):
    """Chapter lines of the table of contents, read from the units before the first chapter: [(roman, title, page)]."""
    lines = []
    for u in units[:before_unit]:
        lines += u["text"].split("\n")
    out = []
    for i, ln in enumerate(lines):
        m = re.match(r"^\s*([IVX]+)\.\s+(.+?)\s*\.{3,}\s*(\d+)?\s*$", ln)
        if not m:
            continue
        page = m.group(3)
        if page is None and i + 1 < len(lines) and re.match(r"^\s*\d+\s*$", lines[i + 1]):
            page = lines[i + 1].strip()
        out.append((chapter_number(m.group(1)), bookio.title_case(m.group(2)), int(page) if page else None))
    return out


def book_pages(units, title):
    """Printed page number of each unit (None if unknown). Printed numbers are read from headers/footers;
    a unit without one takes its neighbor's number plus or minus its distance."""
    printed = []
    for u in units:
        lines = u["text"].split("\n")
        z = pageparts.header_zone(lines, title)
        num = z[2] if z else None
        if num is None:
            f = pageparts.footer_number(lines)
            num = f[1] if f else None
        printed.append(num)
    anchors = [i for i, n in enumerate(printed) if n is not None]
    pages = []
    for i in range(len(units)):
        if printed[i] is not None:
            pages.append(printed[i])
            continue
        near = min(anchors, key=lambda a: (abs(a - i), a > i)) if anchors else None
        guess = printed[near] + (i - near) if near is not None else None
        pages.append(guess if guess and guess >= 1 else None)
    return pages, printed


def run(book_dir, book, pdf):
    """Fills book['chapters']. Returns warnings."""
    units = pdftext.units(pdf)
    heads = find_headings(units)
    warnings = []
    if not heads:
        raise SystemExit("no chapter headings found")
    pages, printed = book_pages(units, book["title"])
    toc = {n: (t, p) for n, t, p in toc_chapters(units, heads[0]["unit"])}
    old = {c["number"]: c for c in book.get("chapters", [])}
    chapters = []
    for k, h in enumerate(heads):
        if k + 1 < len(heads):
            nh = heads[k + 1]
            last_unit = nh["unit"] - 1
            lines = units[nh["unit"]]["text"].split("\n")
            z = pageparts.header_zone(lines, book["title"])
            stray = [l for i, l in enumerate(lines[:nh["line"]]) if l.strip() and not (z and z[0] <= i < z[1])]
            if stray:
                warnings.append(("chapter-starts-mid-page", "chapter %d" % nh["number"],
                                 "text sits above the heading on its page; it is not assigned to the previous chapter"))
        else:
            last_unit = len(units) - 1
        ch = {
            "number": h["number"],
            "title": h["title"],
            "first_unit": h["unit"],
            "last_unit": last_unit,
            "pdf_pages": [units[h["unit"]]["pdf_page"], units[last_unit]["pdf_page"]],
            "book_pages": [pages[h["unit"]], pages[last_unit]],
            "switches": old.get(h["number"], {}).get("switches", {}),
        }
        if h["number"] in toc:
            t_title, t_page = toc[h["number"]]
            ch["toc_page"] = t_page
            if t_page is not None and t_page != ch["book_pages"][0]:
                warnings.append(("toc-mismatch", "chapter %d" % h["number"],
                                 "table of contents says page %s, structure found page %s" % (t_page, ch["book_pages"][0])))
            if pageparts.similar(t_title, h["title"]) < 0.9:
                warnings.append(("toc-mismatch", "chapter %d" % h["number"],
                                 "table of contents title '%s' differs from heading '%s'" % (t_title, h["title"])))
        chapters.append(ch)
    warnings.append(("last-chapter-end", "chapter %d" % chapters[-1]["number"],
                     "no later heading, so the chapter is taken to the end of the book (back matter, if any, is included)"))
    book["chapters"] = chapters
    return warnings