diff options
Diffstat (limited to 'app/converter/extractors.py')
| -rw-r--r-- | app/converter/extractors.py | 328 |
1 files changed, 328 insertions, 0 deletions
diff --git a/app/converter/extractors.py b/app/converter/extractors.py new file mode 100644 index 0000000..a564333 --- /dev/null +++ b/app/converter/extractors.py @@ -0,0 +1,328 @@ +"""Text extraction from book files (TXT, PDF, EPUB) and text/HTML cleaning.""" + +import codecs +import logging +import re +import zipfile +from html import unescape +from pathlib import Path +from typing import List, NamedTuple + +try: + from bs4 import BeautifulSoup + BS4_AVAILABLE = True +except ImportError: + BS4_AVAILABLE = False + +logger = logging.getLogger(__name__) + + +class Section(NamedTuple): + """A titled chunk of a book (e.g. an EPUB chapter).""" + + title: str + text: str + + +class Book(NamedTuple): + """A book's metadata plus its titled sections.""" + + title: str + author: str + sections: List[Section] + + +def extract_text(file_path: Path) -> str: + """Extract text from a book file based on its extension.""" + extension = file_path.suffix.lower() + if extension == ".txt": + return _extract_txt(file_path) + if extension == ".pdf": + return _extract_pdf(file_path) + if extension == ".epub": + return extract_epub(file_path) + raise ValueError(f"Unsupported file format: {extension}") + + +def extract_sections(file_path: Path) -> List[Section]: + """Extract the book's text as titled sections (chapters). + + EPUB files are split on their spine documents so they can be converted + one chapter at a time. TXT and PDF files have no chapter structure and + always yield a single section. + """ + if file_path.suffix.lower() == ".epub": + chapters = _extract_epub_chapters(file_path) + if not chapters: + raise RuntimeError("All EPUB extraction methods failed") + return chapters + + return [Section(file_path.stem, extract_text(file_path))] + + +def extract_book(file_path: Path) -> Book: + """Extract sections plus book-level metadata (title, author). + + EPUB and PDF files carry embedded metadata; missing fields (and TXT + files, which have none) fall back to the file stem for the title and + an empty author. + """ + title, author = "", "" + extension = file_path.suffix.lower() + if extension == ".epub": + title, author = _epub_metadata(file_path) + elif extension == ".pdf": + title, author = _pdf_metadata(file_path) + return Book(title or file_path.stem, author.strip(), extract_sections(file_path)) + + +def _epub_metadata(file_path: Path) -> tuple: + """Return (title, author) from an EPUB's Dublin Core metadata.""" + try: + import ebooklib + from ebooklib import epub + + book = epub.read_epub(str(file_path)) + title = _first_dc_value(book.get_metadata("DC", "title")) + author = _first_dc_value(book.get_metadata("DC", "creator")) + return title, author + except Exception as exc: + logger.warning("Could not read EPUB metadata: %s", exc) + return "", "" + + +def _pdf_metadata(file_path: Path) -> tuple: + """Return (title, author) from a PDF's document info dictionary.""" + try: + from pypdf import PdfReader + + reader = PdfReader(str(file_path)) + info = reader.metadata or {} + title = str(info.get("/Title") or "") + author = str(info.get("/Author") or "") + return title, author + except Exception as exc: + logger.warning("Could not read PDF metadata: %s", exc) + return "", "" + + +def _first_dc_value(entries) -> str: + """First value of an ebooklib DC metadata list: [(value, ...), ...].""" + if not entries: + return "" + value = entries[0][0] + return str(value).strip() if value else "" + + +def _extract_epub_chapters(file_path: Path) -> List[Section]: + """Return one Section per EPUB spine document (chapter), in reading order.""" + import ebooklib + + book = None + for method in (_read_epub_ebooklib, _read_epub_zipfile, _read_epub_manual): + try: + book = method(file_path) + except Exception as exc: + logger.warning("EPUB chapter method %s failed: %s", method.__name__, exc) + continue + if book: + break + + if book is None: + return [] + + chapters = [] + for title, text in book: + cleaned = clean_html(text) + if cleaned.strip(): + chapters.append(Section(title or file_path.stem, cleaned)) + return chapters + + +def _toc_titles(book) -> dict: + """Flatten an ebooklib TOC into a ``{href: title}`` mapping.""" + titles = {} + + def walk(nodes) -> None: + for node in nodes: + if isinstance(node, (tuple, list)): + walk(node[1] if len(node) > 1 else []) + continue + href = getattr(node, "href", None) + title = getattr(node, "title", None) + if href and title: + titles[href.split("#")[0]] = title + + walk(book.toc) + return titles + + +def _read_epub_ebooklib(file_path: Path): + """Read EPUB spine documents as (title, html) pairs via ebooklib.""" + import ebooklib + from ebooklib import epub + + book = epub.read_epub(str(file_path)) + titles = _toc_titles(book) + items = [] + for entry in book.spine: + item_id = entry[0] if isinstance(entry, (tuple, list)) else entry + try: + item = book.get_item_with_id(item_id) + except Exception as exc: + logger.debug("Skipping EPUB spine item %r: %s", item_id, exc) + continue + if not item or item.get_type() != ebooklib.ITEM_DOCUMENT: + continue + if isinstance(item, epub.EpubNav): + continue + content = item.get_body_content() + if content: + if isinstance(content, bytes): + content = content.decode("utf-8", errors="ignore") + title = (titles.get(item.file_name) + or titles.get(item.get_name()) + or getattr(item, "title", None) + or item.get_name()) + items.append((title, str(content))) + return items + + +def _read_epub_zipfile(file_path: Path): + """Read EPUB HTML members as (title, html) pairs, ordered by filename.""" + items = [] + with zipfile.ZipFile(file_path, "r") as epub_zip: + for file_name in sorted(epub_zip.namelist(), key=_natural_key): + if file_name.lower().endswith((".html", ".xhtml", ".htm")): + try: + content = epub_zip.read(file_name).decode("utf-8", errors="ignore") + items.append((Path(file_name).stem, content)) + except Exception as exc: + logger.debug("Skipping EPUB member %r: %s", file_name, exc) + return items + + +def _read_epub_manual(file_path: Path): + """Last-resort read of any markup-looking EPUB member.""" + skipped_extensions = (".jpg", ".jpeg", ".png", ".gif", ".css", ".js") + items = [] + with zipfile.ZipFile(file_path, "r") as epub_zip: + for file_name in sorted(epub_zip.namelist(), key=_natural_key): + if file_name.lower().endswith(skipped_extensions): + continue + try: + content = epub_zip.read(file_name).decode("utf-8", errors="ignore") + if "<" in content and len(content.strip()) > 100: + items.append((Path(file_name).stem, content)) + except Exception as exc: + logger.debug("Skipping EPUB member %r: %s", file_name, exc) + return items + + +def clean_text(text: str) -> str: + """Normalize whitespace and strip standalone page numbers. + + Page numbers are removed only when they appear as a short number alone on + its own line (before whitespace collapsing), so inline numbers like + "42 years", "1,000" or "3.5" are preserved. + """ + if not text: + return "" + # Standalone page numbers (digits alone on a line) must go BEFORE the + # newline-collapsing step below. + text = re.sub(r"(?m)^\s*\d{1,4}\s*$", " ", text) + text = re.sub(r"\s+", " ", text) + return text.strip() + + +def clean_html(html_content: str) -> str: + """Strip markup, scripts and styles from HTML content.""" + if not html_content: + return "" + + if BS4_AVAILABLE: + try: + soup = BeautifulSoup(html_content, "html.parser") + for tag in soup(["script", "style"]): + tag.decompose() + text = soup.get_text(separator=" ") + return re.sub(r"\s+", " ", text).strip() + except Exception as exc: + logger.debug("BeautifulSoup cleaning failed, falling back to regex: %s", exc) + + # Fallback regex cleaning + html_content = re.sub(r"<style[^>]*>.*?</style>", "", html_content, flags=re.DOTALL | re.IGNORECASE) + html_content = re.sub(r"<script[^>]*>.*?</script>", "", html_content, flags=re.DOTALL | re.IGNORECASE) + html_content = re.sub(r"<[^>]+>", " ", html_content) + html_content = unescape(html_content) + html_content = re.sub(r"\s+", " ", html_content) + return html_content.strip() + + +def extract_epub(file_path: Path) -> str: + """Extract the book's text from EPUB, trying several methods in order.""" + chapters = _extract_epub_chapters(file_path) + if not chapters: + raise RuntimeError("All EPUB extraction methods failed") + return "\n\n".join(section.text for section in chapters) + + +def _natural_key(name: str): + """Sort key that orders numeric runs numerically (chapter2 before chapter10).""" + return [int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", name)] + + +def _extract_txt(file_path: Path) -> str: + """Extract from TXT, handling BOMs and common encodings (latin-1 is the catch-all). + + UTF-16 files without a BOM are detected via NUL bytes; otherwise they would + silently decode as NUL-interleaved UTF-8 or cp1252/latin-1 garbage. + """ + data = file_path.read_bytes() + + if data.startswith((codecs.BOM_UTF32_LE, codecs.BOM_UTF32_BE)): + return clean_text(data.decode("utf-32")) + if data.startswith((codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE)): + return clean_text(data.decode("utf-16")) + if data.startswith(codecs.BOM_UTF8): + return clean_text(data.decode("utf-8-sig")) + + # No BOM: UTF-16 without BOM is common on Windows; detect via NUL bytes. + sample = data[:4096] + even_nuls = sum(1 for i, b in enumerate(sample) if b == 0 and i % 2 == 0) + odd_nuls = sum(1 for i, b in enumerate(sample) if b == 0 and i % 2 == 1) + if even_nuls or odd_nuls: + encoding = "utf-16-be" if even_nuls > odd_nuls else "utf-16-le" + return clean_text(data.decode(encoding)) + + for encoding in ("utf-8", "cp1252", "latin-1"): + try: + return clean_text(data.decode(encoding)) + except UnicodeError: + continue + raise ValueError(f"Could not decode text file: {file_path}") + + +def _extract_pdf(file_path: Path) -> str: + """Extract from PDF.""" + from pypdf import PdfReader + + text = "" + with open(file_path, "rb") as file: + pdf_reader = PdfReader(file) + total_pages = len(pdf_reader.pages) + logger.info("PDF has %d pages", total_pages) + + for page_num, page in enumerate(pdf_reader.pages, 1): + try: + page_text = page.extract_text() or "" + if page_text.strip(): + text += f"\n\n{page_text}" + if page_num % 10 == 0: + logger.debug("Extracted %d/%d pages", page_num, total_pages) + except Exception as exc: + logger.warning("Failed to extract page %d: %s", page_num, exc) + + logger.info("Extracted text from %d pages, %d characters total", total_pages, len(text)) + return clean_text(text) |
