aboutsummaryrefslogtreecommitdiff
path: root/app/converter/extractors.py
diff options
context:
space:
mode:
authorhistoria <historiavg@proton.me>2026-08-24 02:59:26 -0400
committerhistoria <historiavg@proton.me>2026-08-24 02:59:26 -0400
commitf00249db9d1ea051d29aa1bcca869fc4b88e83eb (patch)
treea75f076fac1b63e0b4bf2eb8f54affbcc681a891 /app/converter/extractors.py
parent9dd4f9595be3b1d76a3a07dc3eca90cfaf8a3f97 (diff)
downloadtts-audiobook-generator-f00249db9d1ea051d29aa1bcca869fc4b88e83eb.tar.gz
refactor: add app directory, dir structure change
Diffstat (limited to 'app/converter/extractors.py')
-rw-r--r--app/converter/extractors.py328
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)