"""Source documents, pages, and the page windows detection runs over. Every index this project ever reports is an offset into ``Page.text`` as produced here. That makes the extractor part of the contract, so the extractor name and a per-page checksum are recorded with the source: if you change extractors, stored offsets from the old run no longer mean anything, and ``verify_offsets`` says so. """ from __future__ import annotations import hashlib import re from collections.abc import Iterator, Sequence from dataclasses import dataclass, field from pathlib import Path from typing import Any PAGE_SEPARATOR = "\n" def sha256_file(path: Path | str, chunk: int = 1 << 20) -> str: digest = hashlib.sha256() with Path(path).open("rb") as stream: while block := stream.read(chunk): digest.update(block) return digest.hexdigest() def sha1_text(text: str) -> str: return hashlib.sha1(text.encode("utf-8")).hexdigest() _LIGATURES = { "ff": "ff", "fi": "fi", "fl": "fl", "ffi": "ffi", "ffl": "ffl", " ": " ", "’": "'", "‘": "'", "“": '"', "”": '"', "–": "-", "—": "--", "−": "-", } _TRAILING_SPACE = re.compile(r"[ \t]+\n") def normalise_page_text(text: str) -> str: """Light, *stable* cleanup. Deliberately not aggressive: offsets depend on it.""" for bad, good in _LIGATURES.items(): text = text.replace(bad, good) text = text.replace("\r\n", "\n").replace("\r", "\n") text = _TRAILING_SPACE.sub("\n", text) return text class OcrUnavailable(RuntimeError): """Raised when OCR was asked for and could not read a single page.""" @dataclass(frozen=True) class Page: number: int # 1-based, physical page order in the file text: str @property def checksum(self) -> str: return sha1_text(self.text) def __len__(self) -> int: return len(self.text) # -------------------------------------------------------------------------- # extraction def extract_pdf_pages(path: Path | str) -> tuple[list[str], str]: """Return (page texts, extractor id). PyMuPDF if present, else pypdf.""" path = Path(path) backend = None try: # the modern name; `fitz` is the old one and warns on import import pymupdf as backend # type: ignore except ImportError: try: import fitz as backend # type: ignore except ImportError: backend = None if backend is not None: with backend.open(path) as document: pages = [page.get_text("text") for page in document] version = getattr(backend, "__version__", "") or getattr(backend, "VersionBind", "") return pages, f"pymupdf:{version}" if version else "pymupdf" try: from pypdf import PdfReader # type: ignore except ImportError as error: # pragma: no cover - environment problem raise RuntimeError( "No PDF backend available: pip install pymupdf (preferred) or pypdf" ) from error reader = PdfReader(str(path)) return [page.extract_text() or "" for page in reader.pages], "pypdf" def extract_text_pages(path: Path | str, separator: str = "\f") -> tuple[list[str], str]: """Plain-text sources: form feed marks a page break, like ``pdftotext`` emits.""" text = Path(path).read_text(encoding="utf-8", errors="replace") return text.split(separator), "text" @dataclass class Document: """A source book: its pages, and the hash the database keys everything by.""" path: Path sha256: str pages: list[Page] extractor: str title: str = "" meta: dict = field(default_factory=dict) @classmethod def load( cls, path: Path | str, title: str = "", *, mode: str = "text", ocr: Any = None, threshold: float = 0.75, first_page: int | None = None, last_page: int | None = None, progress: Any = None, ) -> Document: """Load a source. ``mode`` decides where the text comes from: ``text`` the embedded text, as the PDF stores it. Fast, free, and wrong whenever the document's maths fonts carry a custom encoding. ``ocr`` render every page and read the pixels. Slow and costed per page, but it sees what a human sees. ``auto`` embedded text, except on pages that fail a quality check -- which is usually exactly the pages with the mathematics on them. ``first_page``/``last_page`` bound which pages are OCR'd, so a 25-page trial of a 600-page book costs 25 calls rather than 600. Pages outside the range keep their embedded text. """ path = Path(path) if not path.is_file(): raise FileNotFoundError(path) if path.suffix.lower() == ".pdf": raw_pages, extractor = extract_pdf_pages(path) else: raw_pages, extractor = extract_text_pages(path) pages = [Page(i + 1, normalise_page_text(text)) for i, text in enumerate(raw_pages)] document = cls( path=path, sha256=sha256_file(path), pages=pages, extractor=extractor, title=title or path.stem, ) if mode != "text" and ocr is not None: document._apply_ocr(ocr, mode, threshold, first_page, last_page, progress) return document def _apply_ocr(self, ocr, mode, threshold, first_page, last_page, progress) -> None: from .ocr import text_quality say = progress or (lambda _message: None) low = max(1, first_page or 1) high = min(len(self.pages), last_page or len(self.pages)) replaced: list[int] = [] skipped: list[int] = [] for index, page in enumerate(self.pages): if not low <= page.number <= high: continue quality = text_quality(page.text) if mode == "auto" and quality >= threshold: continue text = ocr.page_text(self.path, self.sha256, page.number) if text is None: skipped.append(page.number) say(f" page {page.number}: OCR failed, keeping embedded text") continue self.pages[index] = Page(page.number, normalise_page_text(text)) replaced.append(page.number) say(f" page {page.number}: OCR (embedded text scored {quality})") # One page the model chokes on is a bad page; every page failing is a broken # configuration -- a missing key, an unavailable model, no network. Falling # back page by page is right in the first case and disastrous in the second, # because the fallback text is exactly the mangled text OCR exists to replace: # the run completes, costs a full detection pass, and yields nothing. It is # cheaper to stop here than to explain the empty database afterwards. if skipped and not replaced: reason = next(iter(ocr.failures.values()), "no reason recorded") raise OcrUnavailable( f"OCR produced nothing for all {len(skipped)} page(s) via " f"{ocr.backend.name}; first failure: {reason}" ) if replaced: self.extractor = f"{self.extractor}+{ocr.backend.name}@{ocr.dpi}dpi" self.meta["ocr"] = { "mode": mode, "backend": ocr.backend.name, "dpi": ocr.dpi, "threshold": threshold, "pages_replaced": replaced, "pages_failed": skipped, } def __len__(self) -> int: return len(self.pages) def page(self, number: int) -> Page: if not 1 <= number <= len(self.pages): raise IndexError(f"page {number} out of range 1..{len(self.pages)}") return self.pages[number - 1] def page_range(self, first: int | None = None, last: int | None = None) -> list[Page]: first = max(1, first or 1) last = min(len(self.pages), last or len(self.pages)) return [p for p in self.pages if first <= p.number <= last] def slice(self, start_page: int, start_index: int, end_page: int, end_index: int) -> str: """Re-read a stored span straight out of the document.""" if end_page < start_page or (end_page == start_page and end_index < start_index): raise ValueError("span ends before it starts") chunks: list[str] = [] for number in range(start_page, end_page + 1): text = self.page(number).text begin = start_index if number == start_page else 0 finish = end_index if number == end_page else len(text) chunks.append(text[begin:finish]) return PAGE_SEPARATOR.join(chunks) def checksums(self) -> dict[int, str]: return {page.number: page.checksum for page in self.pages} # -------------------------------------------------------------------------- # windows @dataclass class PageWindow: """One or more consecutive pages presented to the model as a single text. ``locate`` maps any offset in that joined text back to (page number, offset in page), which is what the detector reports and the database stores. """ pages: Sequence[Page] def __post_init__(self) -> None: if not self.pages: raise ValueError("a window needs at least one page") parts: list[str] = [] spans: list[tuple[int, int, int]] = [] # (global start, global end, page number) cursor = 0 for position, page in enumerate(self.pages): if position: cursor += len(PAGE_SEPARATOR) spans.append((cursor, cursor + len(page.text), page.number)) parts.append(page.text) cursor += len(page.text) self.text = PAGE_SEPARATOR.join(parts) self._spans = spans @property def first_page(self) -> int: return self.pages[0].number @property def last_page(self) -> int: return self.pages[-1].number def __len__(self) -> int: return len(self.text) def locate(self, index: int) -> tuple[int, int]: """Global offset -> (page number, offset within that page). Clamps politely.""" index = max(0, min(index, len(self.text))) for start, end, number in self._spans: if index < start: # landed inside a separator: belongs to the page ahead return number, 0 if index <= end: return number, index - start start, end, number = self._spans[-1] return number, end - start def globalise(self, page_number: int, offset: int) -> int: for start, end, number in self._spans: if number == page_number: return start + max(0, min(offset, end - start)) raise ValueError(f"page {page_number} is not in this window") def pages_covered(self) -> list[int]: return [page.number for page in self.pages] def windows(pages: Sequence[Page], size: int) -> Iterator[PageWindow]: """Non-overlapping windows of at most ``size`` pages (used by tests/tools).""" for start in range(0, len(pages), size): yield PageWindow(pages[start : start + size])