diff --git a/abogen/text_extractor.py b/abogen/text_extractor.py index 37ad876..e175051 100644 --- a/abogen/text_extractor.py +++ b/abogen/text_extractor.py @@ -1,18 +1,38 @@ from __future__ import annotations +import datetime +import logging +import re +import textwrap +import urllib.parse from dataclasses import dataclass, field from pathlib import Path -import re -from typing import List, Sequence +from typing import Dict, List, Optional, Tuple -import ebooklib -import fitz -import markdown -from bs4 import BeautifulSoup -from ebooklib import epub +import ebooklib # type: ignore[import] +import fitz # type: ignore[import] +import markdown # type: ignore[import] +from bs4 import BeautifulSoup, NavigableString # type: ignore[import] +from ebooklib import epub # type: ignore[import] from .utils import clean_text, detect_encoding +logger = logging.getLogger(__name__) + +METADATA_PATTERN = re.compile(r"<>", re.DOTALL) +CHAPTER_PATTERN = re.compile(r"<>", re.IGNORECASE) +METADATA_KEY_MAP: Dict[str, str] = { + "TITLE": "title", + "ARTIST": "artist", + "ALBUM": "album", + "YEAR": "year", + "ALBUM_ARTIST": "album_artist", + "ALBUMARTIST": "album_artist", + "COMPOSER": "composer", + "GENRE": "genre", + "DATE": "date", +} + @dataclass class ExtractedChapter: @@ -27,7 +47,7 @@ class ExtractedChapter: @dataclass class ExtractionResult: chapters: List[ExtractedChapter] - metadata: dict[str, str] = field(default_factory=dict) + metadata: Dict[str, str] = field(default_factory=dict) @property def combined_text(self) -> str: @@ -38,6 +58,24 @@ class ExtractionResult: return sum(chapter.characters for chapter in self.chapters) +@dataclass +class MetadataSource: + title: Optional[str] = None + authors: List[str] = field(default_factory=list) + description: Optional[str] = None + publisher: Optional[str] = None + publication_year: Optional[str] = None + + +@dataclass +class NavEntry: + src: str + title: str + doc_href: str + position: int + doc_order: int + + def extract_from_path(path: Path) -> ExtractionResult: suffix = path.suffix.lower() if suffix == ".txt": @@ -57,20 +95,27 @@ def _extract_plaintext(path: Path) -> ExtractionResult: return _extract_from_string(raw, default_title=path.stem) -METADATA_PATTERN = re.compile(r"<>", re.DOTALL) -CHAPTER_PATTERN = re.compile(r"<>", re.IGNORECASE) - - def _extract_from_string(raw: str, default_title: str) -> ExtractionResult: - metadata, body = _strip_metadata(raw) + raw_metadata, body = _strip_metadata(raw) chapters = _split_chapters(body, default_title) + normalized_tags = _normalize_metadata_keys(raw_metadata) + chapter_count = len(chapters) + artist_value = normalized_tags.get("artist") + authors = [name.strip() for name in artist_value.split(",") if name.strip()] if artist_value else [] + metadata_source = MetadataSource( + title=normalized_tags.get("title") or default_title, + authors=authors, + publication_year=normalized_tags.get("year"), + ) + metadata = _build_metadata_payload(metadata_source, chapter_count, "text", default_title) + metadata.update(normalized_tags) if not chapters: chapters = [ExtractedChapter(title=default_title, text="")] return ExtractionResult(chapters=chapters, metadata=metadata) -def _strip_metadata(content: str) -> tuple[dict[str, str], str]: - metadata: dict[str, str] = {} +def _strip_metadata(content: str) -> Tuple[Dict[str, str], str]: + metadata: Dict[str, str] = {} def _replacer(match: re.Match) -> str: key = match.group(1).strip().upper() @@ -107,82 +152,710 @@ def _split_chapters(content: str, default_title: str) -> List[ExtractedChapter]: return chapters -def _extract_pdf(path: Path) -> ExtractionResult: - document = fitz.open(str(path)) - chapters: List[ExtractedChapter] = [] - for index, page in enumerate(document): - text = clean_text(page.get_text()) - if not text: +def _normalize_metadata_keys(metadata: Dict[str, str]) -> Dict[str, str]: + normalized: Dict[str, str] = {} + for key, value in metadata.items(): + if not value: continue - title = f"Page {index + 1}" - chapters.append(ExtractedChapter(title=title, text=text)) + mapped = METADATA_KEY_MAP.get(key.upper(), key.lower()) + normalized[mapped] = value + return normalized + + +def _build_metadata_payload( + metadata_source: MetadataSource, + chapter_count: int, + file_type: str, + default_title: str, +) -> Dict[str, str]: + now_year = str(datetime.datetime.now().year) + title = metadata_source.title.strip() if metadata_source.title else default_title + if not title: + title = default_title + authors = [author for author in metadata_source.authors if author.strip()] + if not authors: + authors = ["Unknown"] + authors_text = ", ".join(authors) + if chapter_count <= 0: + chapter_count = 1 + chapter_label = "Chapters" if file_type in {"epub", "markdown"} else "Pages" + metadata = { + "TITLE": title, + "ARTIST": authors_text, + "ALBUM": f"{title} ({chapter_count} {chapter_label})", + "YEAR": metadata_source.publication_year or now_year, + "ALBUM_ARTIST": authors_text, + "COMPOSER": "Narrator", + "GENRE": "Audiobook", + } + return _normalize_metadata_keys(metadata) + + +def _extract_pdf(path: Path) -> ExtractionResult: + metadata_source = MetadataSource() + chapters: List[ExtractedChapter] = [] + with fitz.open(str(path)) as document: + metadata_source = _collect_pdf_metadata(document) + for index, page in enumerate(document): + text = _clean_pdf_text(page.get_text()) + if not text: + continue + title = f"Page {index + 1}" + chapters.append(ExtractedChapter(title=title, text=text)) if not chapters: chapters.append(ExtractedChapter(title=path.stem, text="")) - return ExtractionResult(chapters) + metadata = _build_metadata_payload(metadata_source, len(chapters), "pdf", path.stem) + return ExtractionResult(chapters=chapters, metadata=metadata) + + +def _collect_pdf_metadata(document: fitz.Document) -> MetadataSource: + metadata = MetadataSource() + info = document.metadata or {} + if info.get("title"): + metadata.title = info["title"] + if info.get("author"): + metadata.authors = [info["author"]] + if info.get("subject"): + metadata.description = info["subject"] + if info.get("keywords"): + keywords = info["keywords"] + if metadata.description: + metadata.description = f"{metadata.description}\n\nKeywords: {keywords}" + else: + metadata.description = f"Keywords: {keywords}" + if info.get("creator"): + metadata.publisher = info["creator"] + for key in ("creationDate", "modDate"): + value = info.get(key) + if not value: + continue + match = re.search(r"D:(\d{4})", value) + if match: + metadata.publication_year = match.group(1) + break + return metadata + + +def _clean_pdf_text(text: str) -> str: + cleaned = clean_text(text) + cleaned = re.sub(r"\[\s*\d+\s*\]", "", cleaned) + cleaned = re.sub(r"^\s*\d+\s*$", "", cleaned, flags=re.MULTILINE) + cleaned = re.sub(r"\s+\d+\s*$", "", cleaned, flags=re.MULTILINE) + cleaned = re.sub(r"\s+[-–—]\s*\d+\s*[-–—]?\s*$", "", cleaned, flags=re.MULTILINE) + return cleaned.strip() def _extract_markdown(path: Path) -> ExtractionResult: encoding = detect_encoding(str(path)) raw = path.read_text(encoding=encoding, errors="replace") - html = markdown.markdown(raw, extensions=["toc", "fenced_code"]) - soup = BeautifulSoup(html, "html.parser") - headings = soup.find_all([f"h{i}" for i in range(1, 7)]) - chapters: List[ExtractedChapter] = [] - if headings: - for heading in headings: - sibling_text = _collect_heading_text(heading) - text = clean_text(sibling_text) - if text: - chapters.append(ExtractedChapter(title=heading.get_text(strip=True), text=text)) + metadata_source, chapters = _parse_markdown(raw, path.stem) if not chapters: - chapters.append(ExtractedChapter(title=path.stem, text=clean_text(raw))) - return ExtractionResult(chapters) + chapters = [ExtractedChapter(title=metadata_source.title or path.stem, text=clean_text(raw))] + metadata = _build_metadata_payload(metadata_source, len(chapters), "markdown", path.stem) + return ExtractionResult(chapters=chapters, metadata=metadata) -def _collect_heading_text(node) -> str: - texts: List[str] = [] - for sibling in node.next_siblings: - if getattr(sibling, "name", None) and sibling.name.startswith("h"): - break - text = getattr(sibling, "get_text", lambda **_: "")() - if text: - texts.append(text) - return "\n".join(texts) +def _parse_markdown(raw: str, default_title: str) -> Tuple[MetadataSource, List[ExtractedChapter]]: + metadata = MetadataSource() + text = textwrap.dedent(raw) + frontmatter_match = re.match(r"^---\s*\n(.*?)\n---\s*\n", text, re.DOTALL) + if frontmatter_match: + frontmatter = frontmatter_match.group(1) + _parse_markdown_frontmatter(frontmatter, metadata) + text_body = text[frontmatter_match.end():] + else: + text_body = text + + md = markdown.Markdown(extensions=["toc", "fenced_code"]) + html = md.convert(text_body) + toc_tokens = getattr(md, "toc_tokens", None) or [] + + if not toc_tokens: + cleaned = clean_text(text_body) + title = metadata.title or default_title + chapters = [ExtractedChapter(title=title, text=cleaned)] if cleaned else [] + return metadata, chapters + + headers: List[dict] = [] + + def _flatten_tokens(tokens): + for token in tokens: + headers.append(token) + if token.get("children"): + _flatten_tokens(token["children"]) + + _flatten_tokens(toc_tokens) + + header_positions: List[Tuple[str, int, str]] = [] + for header in headers: + header_id = header.get("id") + if not header_id: + continue + id_pattern = f'id="{header_id}"' + pos = html.find(id_pattern) + if pos == -1: + continue + tag_start = html.rfind("<", 0, pos) + name = str(header.get("name", header_id)) + header_positions.append((header_id, tag_start, name)) + + header_positions.sort(key=lambda item: item[1]) + + chapters: List[ExtractedChapter] = [] + for index, (header_id, start, name) in enumerate(header_positions): + end = header_positions[index + 1][1] if index + 1 < len(header_positions) else len(html) + section_html = html[start:end] + section_soup = BeautifulSoup(section_html, "html.parser") + header_tag = section_soup.find(attrs={"id": header_id}) + if header_tag: + header_tag.decompose() + section_text = clean_text(section_soup.get_text()).strip() + if not section_text: + continue + chapters.append(ExtractedChapter(title=name.strip(), text=section_text)) + + if not metadata.title: + first_h1 = next((header for header in headers if header.get("level") == 1 and header.get("name")), None) + if first_h1: + metadata.title = str(first_h1["name"]) + + return metadata, chapters + + +def _parse_markdown_frontmatter(frontmatter: str, metadata: MetadataSource) -> None: + title_match = re.search(r"^title:\s*(.+)$", frontmatter, re.MULTILINE | re.IGNORECASE) + if title_match: + metadata.title = title_match.group(1).strip().strip('"\'') + + author_match = re.search(r"^author:\s*(.+)$", frontmatter, re.MULTILINE | re.IGNORECASE) + if author_match: + metadata.authors = [author_match.group(1).strip().strip('"\'')] + + desc_match = re.search(r"^description:\s*(.+)$", frontmatter, re.MULTILINE | re.IGNORECASE) + if desc_match: + metadata.description = desc_match.group(1).strip().strip('"\'') + + date_match = re.search(r"^date:\s*(.+)$", frontmatter, re.MULTILINE | re.IGNORECASE) + if date_match: + date_str = date_match.group(1).strip().strip('\"\'') + year_match = re.search(r"\b(19|20)\d{2}\b", date_str) + if year_match: + metadata.publication_year = year_match.group(0) def _extract_epub(path: Path) -> ExtractionResult: - book = epub.read_epub(str(path)) - chapters: List[ExtractedChapter] = [] - spine_docs: Sequence[str] = [item[0] for item in book.spine] - for item in book.get_items_of_type(ebooklib.ITEM_DOCUMENT): - name = item.get_name() - if name not in spine_docs: - continue - html_bytes = item.get_content() - soup = BeautifulSoup(html_bytes, "html.parser") + extractor = EpubExtractor(path) + return extractor.extract() + + +class EpubExtractor: + def __init__(self, path: Path) -> None: + self.path = path + self.book = epub.read_epub(str(path)) + self.doc_content: Dict[str, str] = {} + self.spine_docs: List[str] = [] + + def extract(self) -> ExtractionResult: + metadata_source = self._collect_metadata() + try: + chapters = self._process_nav() + except Exception as exc: + logger.warning( + "EPUB navigation processing failed for %s: %s. Falling back to spine order.", + self.path.name, + exc, + exc_info=True, + ) + chapters = self._process_spine_fallback() + if not chapters: + chapters = [ExtractedChapter(title=self.path.stem, text="")] + metadata = _build_metadata_payload(metadata_source, len(chapters), "epub", self.path.stem) + return ExtractionResult(chapters=chapters, metadata=metadata) + + def _collect_metadata(self) -> MetadataSource: + metadata = MetadataSource() + try: + title_items = self.book.get_metadata("DC", "title") + if title_items: + metadata.title = title_items[0][0] + except Exception as exc: + logger.debug("Failed to extract EPUB title metadata: %s", exc) + + try: + author_items = self.book.get_metadata("DC", "creator") + if author_items: + metadata.authors = [author[0] for author in author_items if author and author[0]] + except Exception as exc: + logger.debug("Failed to extract EPUB author metadata: %s", exc) + + try: + desc_items = self.book.get_metadata("DC", "description") + if desc_items: + metadata.description = desc_items[0][0] + except Exception as exc: + logger.debug("Failed to extract EPUB description metadata: %s", exc) + + try: + publisher_items = self.book.get_metadata("DC", "publisher") + if publisher_items: + metadata.publisher = publisher_items[0][0] + except Exception as exc: + logger.debug("Failed to extract EPUB publisher metadata: %s", exc) + + try: + date_items = self.book.get_metadata("DC", "date") + if date_items: + date_str = date_items[0][0] + year_match = re.search(r"\b(19|20)\d{2}\b", date_str) + metadata.publication_year = year_match.group(0) if year_match else date_str + except Exception as exc: + logger.debug("Failed to extract EPUB publication year metadata: %s", exc) + + return metadata + + def _process_nav(self) -> List[ExtractedChapter]: + nav_item, nav_type = self._find_navigation_item() + if not nav_item or not nav_type: + raise ValueError("No navigation document found") + + parser_type = "html.parser" if nav_type == "html" else "xml" + nav_content = nav_item.get_content().decode("utf-8", errors="ignore") + nav_soup = BeautifulSoup(nav_content, parser_type) + + self.spine_docs = self._build_spine_docs() + doc_order = {href: index for index, href in enumerate(self.spine_docs)} + doc_order_decoded = {urllib.parse.unquote(href): index for href, index in doc_order.items()} + + nav_targets = self._collect_nav_targets(nav_soup, nav_type) + self._cache_relevant_documents(doc_order, nav_targets) + + ordered_entries: List[NavEntry] = [] + if nav_type == "ncx": + nav_map = nav_soup.find("navMap") + if not nav_map: + raise ValueError("NCX navigation missing ") + for nav_point in nav_map.find_all("navPoint", recursive=False): + self._parse_ncx_navpoint(nav_point, ordered_entries, doc_order, doc_order_decoded) + else: + toc_nav = nav_soup.find("nav", attrs={"epub:type": "toc"}) + if toc_nav is None: + for nav in nav_soup.find_all("nav"): + if nav.find("ol"): + toc_nav = nav + break + if toc_nav is None: + raise ValueError("NAV HTML missing TOC structure") + top_ol = toc_nav.find("ol", recursive=False) + if top_ol is None: + raise ValueError("TOC navigation missing
    ") + for li in top_ol.find_all("li", recursive=False): + self._parse_html_nav_li(li, ordered_entries, doc_order, doc_order_decoded) + + if not ordered_entries: + raise ValueError("No navigation entries found") + + ordered_entries.sort(key=lambda entry: (entry.doc_order, entry.position)) + chapters = self._slice_entries(ordered_entries) + self._append_prefix_content(ordered_entries, chapters) + return chapters + + def _process_spine_fallback(self) -> List[ExtractedChapter]: + chapters: List[ExtractedChapter] = [] + self.spine_docs = self._build_spine_docs() + self.doc_content = {} + + for item in self.book.get_items_of_type(ebooklib.ITEM_DOCUMENT): + href = item.get_name() + if href not in self.spine_docs: + continue + try: + html_content = item.get_content().decode("utf-8", errors="ignore") + except Exception as exc: + logger.error("Error decoding EPUB document %s: %s", href, exc) + html_content = "" + self.doc_content[href] = html_content + + for index, doc_href in enumerate(self.spine_docs): + html_content = self.doc_content.get(doc_href, "") + if not html_content: + continue + text = self._html_to_text(html_content) + if not text: + continue + title = self._resolve_document_title(html_content, fallback=f"Untitled Chapter {index + 1}") + chapters.append(ExtractedChapter(title=title, text=text)) + return chapters + + def _find_navigation_item(self) -> Tuple[Optional[ebooklib.epub.EpubItem], Optional[str]]: + nav_item: Optional[ebooklib.epub.EpubItem] = None + nav_type: Optional[str] = None + + nav_items = list(self.book.get_items_of_type(ebooklib.ITEM_NAVIGATION)) + if nav_items: + preferred = next( + ( + item + for item in nav_items + if "nav" in item.get_name().lower() and item.get_name().lower().endswith((".xhtml", ".html")) + ), + None, + ) + if preferred: + nav_item = preferred + nav_type = "html" + else: + html_nav = next( + ( + item + for item in nav_items + if item.get_name().lower().endswith((".xhtml", ".html")) + ), + None, + ) + if html_nav: + nav_item = html_nav + nav_type = "html" + + if not nav_item and nav_items: + ncx_candidate = next( + (item for item in nav_items if item.get_name().lower().endswith(".ncx")), + None, + ) + if ncx_candidate: + nav_item = ncx_candidate + nav_type = "ncx" + + if not nav_item: + ncx_constant = getattr(epub, "ITEM_NCX", None) + if ncx_constant is not None: + ncx_items = list(self.book.get_items_of_type(ncx_constant)) + if ncx_items: + nav_item = ncx_items[0] + nav_type = "ncx" + + if not nav_item: + for item in self.book.get_items_of_type(ebooklib.ITEM_DOCUMENT): + try: + html_content = item.get_content().decode("utf-8", errors="ignore") + except Exception: + continue + if " List[str]: + docs: List[str] = [] + for spine_entry in self.book.spine: + item_id = spine_entry[0] + item = self.book.get_item_with_id(item_id) + if item: + docs.append(item.get_name()) + return docs + + def _collect_nav_targets(self, nav_soup: BeautifulSoup, nav_type: str) -> List[str]: + targets: List[str] = [] + if nav_type == "ncx": + for content_node in nav_soup.find_all("content"): + src = content_node.get("src") + if src: + targets.append(src.split("#", 1)[0]) + else: + for link in nav_soup.find_all("a"): + href = link.get("href") + if href: + targets.append(href.split("#", 1)[0]) + return targets + + def _cache_relevant_documents(self, doc_order: Dict[str, int], nav_targets: List[str]) -> None: + needed: set[str] = set(doc_order.keys()) + for target in nav_targets: + needed.add(target) + needed.add(urllib.parse.unquote(target)) + + self.doc_content = {} + for item in self.book.get_items_of_type(ebooklib.ITEM_DOCUMENT): + href = item.get_name() + if href not in needed and urllib.parse.unquote(href) not in needed: + continue + try: + html_content = item.get_content().decode("utf-8", errors="ignore") + except Exception as exc: + logger.error("Error decoding EPUB document %s: %s", href, exc) + html_content = "" + self.doc_content[href] = html_content + + def _parse_ncx_navpoint( + self, + nav_point, + ordered_entries: List[NavEntry], + doc_order: Dict[str, int], + doc_order_decoded: Dict[str, int], + ) -> None: + nav_label = nav_point.find("navLabel") + content = nav_point.find("content") + title = ( + nav_label.find("text").get_text(strip=True) + if nav_label and nav_label.find("text") + else "Untitled Section" + ) + src = content.get("src") if content and content.has_attr("src") else None + + if src: + base_href, fragment = src.split("#", 1) if "#" in src else (src, None) + doc_key, doc_idx = self._find_doc_key(base_href, doc_order, doc_order_decoded) + if doc_key is not None and doc_idx is not None: + position = self._find_position_robust(doc_key, fragment) + ordered_entries.append( + NavEntry( + src=src, + title=title, + doc_href=doc_key, + position=position, + doc_order=doc_idx, + ) + ) + else: + logger.warning( + "Navigation entry '%s' points to '%s', which is not in the spine.", + title, + base_href, + ) + + for child_navpoint in nav_point.find_all("navPoint", recursive=False): + self._parse_ncx_navpoint(child_navpoint, ordered_entries, doc_order, doc_order_decoded) + + def _parse_html_nav_li( + self, + li_element, + ordered_entries: List[NavEntry], + doc_order: Dict[str, int], + doc_order_decoded: Dict[str, int], + ) -> None: + link = li_element.find("a", recursive=False) + span_text = li_element.find("span", recursive=False) + title = "Untitled Section" + + if link and link.has_attr("href"): + src = link["href"] + title = link.get_text(strip=True) or title + else: + src = None + if span_text: + title = span_text.get_text(strip=True) or title + else: + text = "".join(t for t in li_element.stripped_strings) + if text: + title = text + + title = title.strip() or "Untitled Section" + + if src: + base_href, fragment = src.split("#", 1) if "#" in src else (src, None) + doc_key, doc_idx = self._find_doc_key(base_href, doc_order, doc_order_decoded) + if doc_key is not None and doc_idx is not None: + position = self._find_position_robust(doc_key, fragment) + ordered_entries.append( + NavEntry( + src=src, + title=title, + doc_href=doc_key, + position=position, + doc_order=doc_idx, + ) + ) + else: + logger.warning( + "Navigation entry '%s' points to '%s', which is not in the spine.", + title, + base_href, + ) + + for child_ol in li_element.find_all("ol", recursive=False): + for child_li in child_ol.find_all("li", recursive=False): + self._parse_html_nav_li(child_li, ordered_entries, doc_order, doc_order_decoded) + + def _find_doc_key( + self, + base_href: str, + doc_order: Dict[str, int], + doc_order_decoded: Dict[str, int], + ) -> Tuple[Optional[str], Optional[int]]: + candidates = {base_href, urllib.parse.unquote(base_href)} + base_name = urllib.parse.unquote(base_href).split("/")[-1].lower() + for key in list(doc_order.keys()) + list(doc_order_decoded.keys()): + if key.split("/")[-1].lower() == base_name: + candidates.add(key) + for candidate in candidates: + if candidate in doc_order: + return candidate, doc_order[candidate] + if candidate in doc_order_decoded: + return candidate, doc_order_decoded[candidate] + return None, None + + def _find_position_robust(self, doc_href: str, fragment_id: Optional[str]) -> int: + if doc_href not in self.doc_content: + logger.warning("Document '%s' not found in cached EPUB content.", doc_href) + return 0 + html_content = self.doc_content[doc_href] + if not fragment_id: + return 0 + + try: + temp_soup = BeautifulSoup(f"
    {html_content}
    ", "html.parser") + target_element = temp_soup.find(id=fragment_id) + if target_element: + tag_str = str(target_element) + pos = html_content.find(tag_str[: min(len(tag_str), 200)]) + if pos != -1: + return pos + except Exception: + logger.debug("BeautifulSoup failed to locate id '%s' in %s", fragment_id, doc_href) + + safe_fragment_id = re.escape(fragment_id) + id_name_pattern = re.compile( + f"<[^>]+(?:id|name)\\s*=\\s*[\"']{safe_fragment_id}[\"']", + re.IGNORECASE, + ) + match = id_name_pattern.search(html_content) + if match: + return match.start() + + id_pos = html_content.find(f'id="{fragment_id}"') + name_pos = html_content.find(f'name="{fragment_id}"') + candidates = [pos for pos in (id_pos, name_pos) if pos != -1] + if candidates: + pos = min(candidates) + tag_start = html_content.rfind("<", 0, pos) + return tag_start if tag_start != -1 else pos + + logger.warning("Anchor '%s' not found in %s. Defaulting to start.", fragment_id, doc_href) + return 0 + + def _slice_entries(self, ordered_entries: List[NavEntry]) -> List[ExtractedChapter]: + chapters: List[ExtractedChapter] = [] + for index, entry in enumerate(ordered_entries): + next_entry = ordered_entries[index + 1] if index + 1 < len(ordered_entries) else None + slice_html = self._slice_entry(entry, next_entry) + text = self._html_to_text(slice_html) + if not text: + continue + title = entry.title or "Untitled Section" + chapters.append(ExtractedChapter(title=title, text=text)) + return chapters + + def _slice_entry( + self, + current_entry: NavEntry, + next_entry: Optional[NavEntry], + ) -> str: + current_doc = current_entry.doc_href + current_pos = current_entry.position + current_html = self.doc_content.get(current_doc, "") + if not current_html: + return "" + + if next_entry and next_entry.doc_href == current_doc: + return current_html[current_pos : next_entry.position] + + slice_html = current_html[current_pos:] + if next_entry: + docs_between = self._docs_between(current_doc, next_entry.doc_href) + for doc_href in docs_between: + slice_html += self.doc_content.get(doc_href, "") + next_doc_html = self.doc_content.get(next_entry.doc_href, "") + slice_html += next_doc_html[: next_entry.position] + else: + for doc_href in self._docs_between(current_doc, None): + slice_html += self.doc_content.get(doc_href, "") + + if not slice_html.strip(): + logger.warning( + "No content found for navigation source '%s'. Using full document fallback.", + current_entry.src, + ) + return current_html + return slice_html + + def _docs_between(self, current_doc: str, next_doc: Optional[str]) -> List[str]: + docs: List[str] = [] + try: + current_idx = self.spine_docs.index(current_doc) + except ValueError: + return docs + + if next_doc is None: + docs.extend(self.spine_docs[current_idx + 1 :]) + return docs + + try: + next_idx = self.spine_docs.index(next_doc) + except ValueError: + return docs + + if current_idx < next_idx: + docs.extend(self.spine_docs[current_idx + 1 : next_idx]) + elif current_idx > next_idx: + docs.extend(self.spine_docs[current_idx + 1 :]) + docs.extend(self.spine_docs[:next_idx]) + return docs + + def _append_prefix_content( + self, + ordered_entries: List[NavEntry], + chapters: List[ExtractedChapter], + ) -> None: + if not ordered_entries: + return + first_entry = ordered_entries[0] + first_doc = first_entry.doc_href + first_pos = first_entry.position + if first_pos <= 0: + return + + prefix_html = "" + try: + first_idx = self.spine_docs.index(first_doc) + except ValueError: + first_idx = -1 + + if first_idx > 0: + for doc_href in self.spine_docs[:first_idx]: + prefix_html += self.doc_content.get(doc_href, "") + prefix_html += self.doc_content.get(first_doc, "")[:first_pos] + prefix_text = self._html_to_text(prefix_html) + if prefix_text and (not chapters or prefix_text != chapters[0].text): + chapters.insert(0, ExtractedChapter(title="Introduction", text=prefix_text)) + + def _html_to_text(self, html: str) -> str: + if not html: + return "" + soup = BeautifulSoup(html, "html.parser") + for tag in soup.find_all(["p", "div"]): + tag.append("\n\n") for ol in soup.find_all("ol"): start = int(ol.get("start", 1)) for idx, li in enumerate(ol.find_all("li", recursive=False)): - number = f"{start + idx}. " + number_text = f"{start + idx}) " if li.string: - li.string.replace_with(number + li.string) + li.string.replace_with(number_text + li.string) else: - li.insert(0, number) + li.insert(0, NavigableString(number_text)) + for tag in soup.find_all(["sup", "sub"]): + tag.decompose() text = clean_text(soup.get_text()) - if not text: - continue - title = _resolve_epub_title(soup, name) - chapters.append(ExtractedChapter(title=title, text=text)) - if not chapters: - chapters.append(ExtractedChapter(title=path.stem, text="")) - return ExtractionResult(chapters) + return text.strip() - -def _resolve_epub_title(soup: BeautifulSoup, fallback: str) -> str: - if soup.title and soup.title.string: - return soup.title.string.strip() - for heading_tag in ("h1", "h2", "h3"): - heading = soup.find(heading_tag) - if heading and heading.get_text(strip=True): - return heading.get_text(strip=True) - return fallback + def _resolve_document_title(self, html_content: str, fallback: str) -> str: + soup = BeautifulSoup(html_content, "html.parser") + if soup.title and soup.title.string: + return soup.title.string.strip() + for heading_tag in ("h1", "h2", "h3"): + heading = soup.find(heading_tag) + if heading and heading.get_text(strip=True): + return heading.get_text(strip=True) + return fallback diff --git a/abogen/web/conversion_runner.py b/abogen/web/conversion_runner.py index 37bb02e..e8a4cad 100644 --- a/abogen/web/conversion_runner.py +++ b/abogen/web/conversion_runner.py @@ -9,7 +9,7 @@ import sys from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path -from typing import Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional import numpy as np import soundfile as sf @@ -43,6 +43,126 @@ class AudioSink: write: Callable[[np.ndarray], None] +def _coerce_truthy(value: Any, default: bool = True) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"true", "1", "yes", "on"}: + return True + if lowered in {"false", "0", "no", "off"}: + return False + return default + if value is None: + return default + return bool(value) + + +def _apply_chapter_overrides( + extracted: List[ExtractedChapter], + overrides: List[Dict[str, Any]], +) -> tuple[List[ExtractedChapter], Dict[str, str], List[str]]: + if not overrides: + return [], {}, [] + + selected: List[ExtractedChapter] = [] + metadata_updates: Dict[str, str] = {} + diagnostics: List[str] = [] + + for position, payload in enumerate(overrides): + if not isinstance(payload, dict): + diagnostics.append( + f"Skipped chapter override at position {position + 1}: unsupported payload type {type(payload).__name__}." + ) + continue + + enabled = _coerce_truthy(payload.get("enabled", True)) + payload["enabled"] = enabled + if not enabled: + continue + + metadata_payload = payload.get("metadata") or {} + if isinstance(metadata_payload, dict): + for key, value in metadata_payload.items(): + if value is None: + continue + metadata_updates[str(key)] = str(value) + + base: Optional[ExtractedChapter] = None + idx_candidate = payload.get("index") + idx_normalized: Optional[int] = None + if isinstance(idx_candidate, int): + idx_normalized = idx_candidate + elif isinstance(idx_candidate, str): + try: + idx_normalized = int(idx_candidate) + except ValueError: + idx_normalized = None + if idx_normalized is not None and 0 <= idx_normalized < len(extracted): + base = extracted[idx_normalized] + payload["index"] = idx_normalized + + if base is None: + source_title = payload.get("source_title") + if isinstance(source_title, str): + base = next((chapter for chapter in extracted if chapter.title == source_title), None) + + if base is None: + candidate_title = payload.get("title") + if isinstance(candidate_title, str): + base = next((chapter for chapter in extracted if chapter.title == candidate_title), None) + + text_override = payload.get("text") + if text_override is not None: + text_value = str(text_override) + elif base is not None: + text_value = base.text + else: + diagnostics.append( + f"Skipped chapter override at position {position + 1}: no text provided and no matching source chapter found." + ) + continue + + title_override = payload.get("title") + if title_override is not None: + title_value = str(title_override) + elif base is not None: + title_value = base.title + else: + title_value = f"Chapter {position + 1}" + + if base and not payload.get("source_title"): + payload["source_title"] = base.title + + payload["title"] = title_value + payload["text"] = text_value + payload["characters"] = len(text_value) + payload.setdefault("order", payload.get("order", position)) + + selected.append(ExtractedChapter(title=title_value, text=text_value)) + + return selected, metadata_updates, diagnostics + + +def _merge_metadata( + extracted: Optional[Dict[str, str]], + overrides: Dict[str, Any], +) -> Dict[str, str]: + merged: Dict[str, str] = {} + if extracted: + for key, value in extracted.items(): + if value is None: + continue + merged[str(key)] = str(value) + for key, value in (overrides or {}).items(): + key_str = str(key) + if value is None: + merged.pop(key_str, None) + else: + merged[key_str] = str(value) + return merged + + def run_conversion_job(job: Job) -> None: job.add_log("Preparing conversion pipeline") canceller = _make_canceller(job) @@ -53,11 +173,29 @@ def run_conversion_job(job: Job) -> None: try: pipeline = _load_pipeline(job) extraction = extract_from_path(job.stored_path) - job.metadata_tags = extraction.metadata or {} + + metadata_overrides: Dict[str, Any] = dict(job.metadata_tags or {}) + if job.chapters: + selected_chapters, chapter_metadata, diagnostics = _apply_chapter_overrides( + extraction.chapters, + job.chapters, + ) + for message in diagnostics: + job.add_log(message, level="warning") + if selected_chapters: + extraction.chapters = selected_chapters + metadata_overrides.update(chapter_metadata) + job.add_log( + f"Chapter overrides applied: {len(selected_chapters)} selected.", + level="info", + ) + else: + raise ValueError("No chapters were enabled in the requested job.") + + job.metadata_tags = _merge_metadata(extraction.metadata, metadata_overrides) total_characters = extraction.total_characters or calculate_text_length(extraction.combined_text) - if job.total_characters == 0: - job.total_characters = total_characters + job.total_characters = total_characters job.add_log(f"Total characters: {job.total_characters:,}") _apply_newline_policy(extraction.chapters, job.replace_single_newlines) @@ -237,9 +375,9 @@ def _select_device() -> str: def _prepare_output_dir(job: Job) -> Path: - from platformdirs import user_desktop_dir + from platformdirs import user_desktop_dir # type: ignore[import-not-found] - default_output = Path(get_user_cache_path("outputs")) + default_output = Path(str(get_user_cache_path("outputs"))) if job.save_mode == "Save to Desktop": directory = Path(user_desktop_dir()) elif job.save_mode == "Save next to input file": diff --git a/abogen/web/service.py b/abogen/web/service.py index f8c097d..ac06705 100644 --- a/abogen/web/service.py +++ b/abogen/web/service.py @@ -6,7 +6,7 @@ import uuid from dataclasses import dataclass, field from enum import Enum from pathlib import Path -from typing import Callable, Dict, Iterable, List, Optional +from typing import Any, Callable, Dict, Iterable, List, Optional, Mapping class JobStatus(str, Enum): @@ -64,7 +64,7 @@ class Job: logs: List[JobLog] = field(default_factory=list) error: Optional[str] = None result: JobResult = field(default_factory=JobResult) - chapters: List[str] = field(default_factory=list) + chapters: List[Dict[str, Any]] = field(default_factory=list) queue_position: Optional[int] = None cancel_requested: bool = False @@ -100,6 +100,18 @@ class Job: "voice_profile": self.voice_profile, "max_subtitle_words": self.max_subtitle_words, }, + "metadata_tags": dict(self.metadata_tags), + "chapters": [ + { + "id": entry.get("id"), + "index": entry.get("index"), + "order": entry.get("order"), + "title": entry.get("title"), + "enabled": bool(entry.get("enabled", True)), + "characters": len(str(entry.get("text", ""))), + } + for entry in self.chapters + ], } @@ -149,7 +161,7 @@ class ConversionService: replace_single_newlines: bool, subtitle_format: str, total_characters: int, - chapters: Optional[Iterable[str]] = None, + chapters: Optional[Iterable[Any]] = None, save_chapters_separately: bool = False, merge_chapters_at_end: bool = True, separate_chapters_format: str = "wav", @@ -157,8 +169,13 @@ class ConversionService: save_as_project: bool = False, voice_profile: Optional[str] = None, max_subtitle_words: int = 50, + metadata_tags: Optional[Mapping[str, Any]] = None, ) -> Job: job_id = uuid.uuid4().hex + normalized_metadata = self._normalize_metadata_tags(metadata_tags) + normalized_chapters = self._normalize_chapters(chapters) + if total_characters <= 0 and normalized_chapters: + total_characters = sum(len(str(entry.get("text", ""))) for entry in normalized_chapters) job = Job( id=job_id, original_filename=original_filename, @@ -180,9 +197,10 @@ class ConversionService: save_as_project=save_as_project, voice_profile=voice_profile, max_subtitle_words=max_subtitle_words, + metadata_tags=normalized_metadata, created_at=time.time(), total_characters=total_characters, - chapters=list(chapters or []), + chapters=normalized_chapters, ) with self._lock: self._jobs[job_id] = job @@ -322,6 +340,128 @@ class ConversionService: if job: job.queue_position = index + @staticmethod + def _coerce_bool(value: Any, default: bool = True) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"true", "1", "yes", "on"}: + return True + if lowered in {"false", "0", "no", "off"}: + return False + return default + if value is None: + return default + return bool(value) + + @staticmethod + def _coerce_optional_int(value: Any) -> Optional[int]: + if value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + @staticmethod + def _normalize_metadata_tags(values: Optional[Mapping[str, Any]]) -> Dict[str, str]: + if not values: + return {} + normalized: Dict[str, str] = {} + for key, raw_value in values.items(): + if raw_value is None: + continue + key_str = str(key).strip() + if not key_str: + continue + normalized[key_str] = str(raw_value) + return normalized + + @classmethod + def _normalize_chapters(cls, chapters: Optional[Iterable[Any]]) -> List[Dict[str, Any]]: + if not chapters: + return [] + + normalized: List[Dict[str, Any]] = [] + for order, raw in enumerate(chapters): + if raw is None: + continue + + if isinstance(raw, str): + raw_dict: Dict[str, Any] = {"title": raw} + elif isinstance(raw, dict): + raw_dict = dict(raw) + else: + continue + + entry: Dict[str, Any] = {} + + id_value = raw_dict.get("id") or raw_dict.get("chapter_id") or raw_dict.get("key") + if id_value is not None: + entry["id"] = str(id_value) + + index_value = ( + cls._coerce_optional_int(raw_dict.get("index")) + or cls._coerce_optional_int(raw_dict.get("original_index")) + or cls._coerce_optional_int(raw_dict.get("source_index")) + or cls._coerce_optional_int(raw_dict.get("chapter_index")) + ) + if index_value is not None: + entry["index"] = index_value + + order_value = ( + cls._coerce_optional_int(raw_dict.get("order")) + or cls._coerce_optional_int(raw_dict.get("position")) + or cls._coerce_optional_int(raw_dict.get("sort")) + or cls._coerce_optional_int(raw_dict.get("sort_order")) + ) + entry["order"] = order_value if order_value is not None else order + + source_title = ( + raw_dict.get("source_title") + or raw_dict.get("original_title") + or raw_dict.get("base_title") + ) + if source_title: + entry["source_title"] = str(source_title) + + title_value = ( + raw_dict.get("title") + or raw_dict.get("name") + or raw_dict.get("label") + or raw_dict.get("chapter") + ) + if title_value is not None: + entry["title"] = str(title_value) + elif source_title: + entry["title"] = str(source_title) + else: + entry["title"] = f"Chapter {order + 1}" + + text_value = raw_dict.get("text") + if text_value is None: + text_value = raw_dict.get("content") or raw_dict.get("body") or raw_dict.get("value") + if text_value is not None: + entry["text"] = str(text_value) + + enabled = cls._coerce_bool( + raw_dict.get("enabled", raw_dict.get("include", raw_dict.get("selected", True))), + True, + ) + if "disabled" in raw_dict and cls._coerce_bool(raw_dict.get("disabled"), False): + enabled = False + entry["enabled"] = enabled + + metadata_payload = raw_dict.get("metadata") or raw_dict.get("metadata_tags") + normalized_metadata = cls._normalize_metadata_tags(metadata_payload) + if normalized_metadata: + entry["metadata"] = normalized_metadata + + normalized.append(entry) + + return normalized + def default_storage_root() -> Path: base = Path.cwd() diff --git a/tests/test_chapter_overrides.py b/tests/test_chapter_overrides.py new file mode 100644 index 0000000..814d44b --- /dev/null +++ b/tests/test_chapter_overrides.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +import sys +import types + + +def _install_dependency_stubs() -> None: + if "ebooklib" not in sys.modules: + ebooklib_stub = types.ModuleType("ebooklib") + epub_stub = types.ModuleType("ebooklib.epub") + setattr(ebooklib_stub, "epub", epub_stub) + sys.modules["ebooklib"] = ebooklib_stub + sys.modules["ebooklib.epub"] = epub_stub + + if "dotenv" not in sys.modules: + dotenv_stub = types.ModuleType("dotenv") + + def _noop(*_, **__): + return None + + setattr(dotenv_stub, "load_dotenv", _noop) + setattr(dotenv_stub, "find_dotenv", lambda *_, **__: "") + sys.modules["dotenv"] = dotenv_stub + + if "numpy" not in sys.modules: + numpy_stub = types.ModuleType("numpy") + + class _DummyArray(list): + pass + + def _zeros(shape, dtype=None): + size = 1 + if isinstance(shape, int): + size = shape + elif shape: + size = 1 + for dimension in shape: + size *= int(dimension) + return [0.0] * size + + setattr(numpy_stub, "ndarray", _DummyArray) + setattr(numpy_stub, "zeros", _zeros) + setattr(numpy_stub, "float32", "float32") + setattr(numpy_stub, "array", lambda data, dtype=None: data) + setattr(numpy_stub, "asarray", lambda data, dtype=None: data) + setattr(numpy_stub, "concatenate", lambda seq, axis=0: sum((list(item) for item in seq), [])) + sys.modules["numpy"] = numpy_stub + + if "soundfile" not in sys.modules: + soundfile_stub = types.ModuleType("soundfile") + + class _DummySoundFile: + def __init__(self, *_, **__): + pass + + def write(self, *_args, **_kwargs): + return None + + def close(self): + return None + + setattr(soundfile_stub, "SoundFile", _DummySoundFile) + setattr(soundfile_stub, "write", lambda *_args, **_kwargs: None) + sys.modules["soundfile"] = soundfile_stub + + if "fitz" not in sys.modules: + sys.modules["fitz"] = types.ModuleType("fitz") + + if "markdown" not in sys.modules: + markdown_stub = types.ModuleType("markdown") + + class _DummyMarkdown: + def __init__(self, *_, **__): + pass + + def convert(self, text: str) -> str: + return text + + setattr(markdown_stub, "Markdown", _DummyMarkdown) + sys.modules["markdown"] = markdown_stub + + if "bs4" not in sys.modules: + bs4_stub = types.ModuleType("bs4") + + class _DummySoup: + def __init__(self, *_, **__): + pass + + def select(self, *_, **__): + return [] + + def find_all(self, *_, **__): + return [] + + setattr(bs4_stub, "BeautifulSoup", _DummySoup) + setattr(bs4_stub, "NavigableString", str) + sys.modules["bs4"] = bs4_stub + + +_install_dependency_stubs() + +from abogen.text_extractor import ExtractedChapter +from abogen.web.conversion_runner import _apply_chapter_overrides, _merge_metadata + + +def _sample_chapters() -> list[ExtractedChapter]: + return [ + ExtractedChapter(title="Chapter 1", text="Original one"), + ExtractedChapter(title="Chapter 2", text="Original two"), + ExtractedChapter(title="Chapter 3", text="Original three"), + ] + + +def test_apply_chapter_overrides_with_custom_text() -> None: + overrides = [ + {"index": 0, "enabled": True, "title": "Intro", "text": "Hello world"}, + {"index": 1, "enabled": False}, + ] + + selected, metadata, diagnostics = _apply_chapter_overrides(_sample_chapters(), overrides) + + assert len(selected) == 1 + assert selected[0].title == "Intro" + assert selected[0].text == "Hello world" + assert overrides[0]["characters"] == len("Hello world") + assert metadata == {} + assert diagnostics == [] + + +def test_apply_chapter_overrides_uses_original_content_when_text_missing() -> None: + overrides = [ + {"index": 1, "enabled": True}, + ] + + selected, metadata, diagnostics = _apply_chapter_overrides(_sample_chapters(), overrides) + + assert len(selected) == 1 + assert selected[0].title == "Chapter 2" + assert selected[0].text == "Original two" + assert overrides[0]["text"] == "Original two" + assert overrides[0]["characters"] == len("Original two") + assert metadata == {} + assert diagnostics == [] + + +def test_apply_chapter_overrides_collects_metadata_updates() -> None: + overrides = [ + { + "index": 2, + "enabled": True, + "metadata": {"artist": "Test Author", "year": 2024}, + } + ] + + selected, metadata, diagnostics = _apply_chapter_overrides(_sample_chapters(), overrides) + + assert len(selected) == 1 + assert metadata == {"artist": "Test Author", "year": "2024"} + assert diagnostics == [] + + +def test_apply_chapter_overrides_reports_diagnostics_for_invalid_payload() -> None: + overrides = [ + {"enabled": True, "title": "Missing"}, + ] + + selected, metadata, diagnostics = _apply_chapter_overrides(_sample_chapters(), overrides) + + assert selected == [] + assert metadata == {} + assert diagnostics and "Skipped chapter override" in diagnostics[0] + + +def test_merge_metadata_prefers_overrides_and_drops_none_values() -> None: + extracted = {"title": "Original", "artist": "Someone"} + overrides = {"artist": "Another", "genre": "Fiction", "year": None} + + merged = _merge_metadata(extracted, overrides) + + assert merged["title"] == "Original" + assert merged["artist"] == "Another" + assert merged["genre"] == "Fiction" + assert "year" not in merged