216 lines
7.4 KiB
Python
216 lines
7.4 KiB
Python
import zipfile
|
||
import re
|
||
from lxml import etree
|
||
from dataclasses import dataclass, field
|
||
|
||
|
||
@dataclass
|
||
class Heading:
|
||
title: str
|
||
level: int
|
||
href: str
|
||
cfi: str = ""
|
||
children: list["Heading"] = field(default_factory=list)
|
||
|
||
|
||
NS = {
|
||
"opf": "http://www.idpf.org/2007/opf",
|
||
"ncx": "http://www.daisy.org/z3986/2005/ncx/",
|
||
"xhtml": "http://www.w3.org/1999/xhtml",
|
||
}
|
||
|
||
|
||
def parse_toc_tree(epub_path: str) -> list[Heading]:
|
||
"""从 EPUB 中提取完整目录树，4 级降级：nav.xhtml -> ncx -> spine -> h 标签"""
|
||
with zipfile.ZipFile(epub_path) as zf:
|
||
opf_path, opf_root = _read_opf(zf)
|
||
opf_dir = _dirname(opf_path)
|
||
|
||
headings = _try_nav(zf, opf_root, opf_dir)
|
||
if headings:
|
||
return headings
|
||
headings = _try_ncx(zf, opf_root, opf_dir)
|
||
if headings:
|
||
return headings
|
||
headings = _try_spine(zf, opf_root, opf_dir)
|
||
if headings:
|
||
return headings
|
||
return _fallback_spine(zf, opf_root, opf_dir)
|
||
|
||
|
||
def _read_opf(zf: zipfile.ZipFile) -> tuple[str, etree._Element]:
|
||
container = etree.fromstring(zf.read("META-INF/container.xml"))
|
||
rootfile = container.find(".//{urn:oasis:names:tc:opendocument:xmlns:container}rootfile")
|
||
opf_path = rootfile.attrib["full-path"] if rootfile is not None else "OEBPS/content.opf"
|
||
opf_root = etree.fromstring(zf.read(opf_path))
|
||
return opf_path, opf_root
|
||
|
||
|
||
def _try_nav(zf: zipfile.ZipFile, opf: etree._Element, opf_dir: str) -> list[Heading] | None:
|
||
"""EPUB3: nav.xhtml（最优先，有原生层级）"""
|
||
manifest = {item.attrib.get("id"): item.attrib.get("href", "") for item in opf.findall("opf:manifest/opf:item", NS)}
|
||
spine_ids = [itemref.attrib.get("idref") for itemref in opf.findall("opf:spine/opf:itemref", NS)]
|
||
|
||
for idref in spine_ids:
|
||
href = manifest.get(idref, "")
|
||
if "nav" in href.lower() or "toc" in href.lower():
|
||
nav_path = _resolve_path(opf_dir, href)
|
||
try:
|
||
return _parse_nav_xhtml(zf, nav_path, opf_dir)
|
||
except Exception:
|
||
continue
|
||
return None
|
||
|
||
|
||
def _parse_nav_xhtml(zf: zipfile.ZipFile, nav_path: str, opf_dir: str) -> list[Heading]:
|
||
html = etree.fromstring(zf.read(nav_path))
|
||
ol = html.find(".//{http://www.w3.org/1999/xhtml}nav[@{http://www.w3.org/1999/xhtml}type='toc']//{http://www.w3.org/1999/xhtml}ol")
|
||
if ol is None:
|
||
ol = html.find(".//{http://www.w3.org/1999/xhtml}nav//{http://www.w3.org/1999/xhtml}ol")
|
||
if ol is None:
|
||
return []
|
||
return _parse_ol(ol, opf_dir, 1)
|
||
|
||
|
||
def _parse_ol(ol, opf_dir: str, base_level: int) -> list[Heading]:
|
||
result = []
|
||
for li in ol.findall("{http://www.w3.org/1999/xhtml}li"):
|
||
a = li.find("{http://www.w3.org/1999/xhtml}a")
|
||
if a is None:
|
||
continue
|
||
title = (a.text or "").strip()
|
||
href = a.attrib.get("href", "")
|
||
heading = Heading(title=title, level=base_level, href=href)
|
||
|
||
sub_ol = li.find("{http://www.w3.org/1999/xhtml}ol")
|
||
if sub_ol is not None:
|
||
heading.children = _parse_ol(sub_ol, opf_dir, base_level + 1)
|
||
|
||
result.append(heading)
|
||
return result
|
||
|
||
|
||
def _try_ncx(zf: zipfile.ZipFile, opf: etree._Element, opf_dir: str) -> list[Heading] | None:
|
||
"""EPUB2: toc.ncx"""
|
||
spine = opf.find("opf:spine", NS)
|
||
if spine is None:
|
||
return None
|
||
ncx_id = spine.attrib.get("toc", "")
|
||
if not ncx_id:
|
||
return None
|
||
|
||
manifest = {item.attrib.get("id"): item.attrib.get("href", "") for item in opf.findall("opf:manifest/opf:item", NS)}
|
||
ncx_href = manifest.get(ncx_id)
|
||
if not ncx_href:
|
||
return None
|
||
|
||
ncx_path = _resolve_path(opf_dir, ncx_href)
|
||
try:
|
||
ncx_root = etree.fromstring(zf.read(ncx_path))
|
||
return _parse_ncx_navmap(ncx_root, opf_dir)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _parse_ncx_navmap(ncx: etree._Element, opf_dir: str) -> list[Heading]:
|
||
navmap = ncx.find("ncx:navMap", NS)
|
||
if navmap is None:
|
||
return []
|
||
return _parse_ncx_points(navmap, opf_dir, 1)
|
||
|
||
|
||
def _parse_ncx_points(parent, opf_dir: str, base_level: int) -> list[Heading]:
|
||
result = []
|
||
for np in parent.findall("ncx:navPoint", NS):
|
||
label = np.find("ncx:navLabel/ncx:text", NS)
|
||
content = np.find("ncx:content", NS)
|
||
title = _text(label) if label is not None else ""
|
||
href = content.attrib.get("src", "") if content is not None else ""
|
||
heading = Heading(title=title, level=base_level, href=href)
|
||
heading.children = _parse_ncx_points(np, opf_dir, base_level + 1)
|
||
result.append(heading)
|
||
return result
|
||
|
||
|
||
def _try_spine(zf: zipfile.ZipFile, opf: etree._Element, opf_dir: str) -> list[Heading] | None:
|
||
"""退化方案：每个 spine 文件当作一章"""
|
||
manifest = {item.attrib.get("id"): item.attrib.get("href", "") for item in opf.findall("opf:manifest/opf:item", NS)}
|
||
spine_ids = [itemref.attrib.get("idref") for itemref in opf.findall("opf:spine/opf:itemref", NS)]
|
||
headings = []
|
||
for i, idref in enumerate(spine_ids):
|
||
href = manifest.get(idref, "")
|
||
if not href:
|
||
continue
|
||
try:
|
||
path = _resolve_path(opf_dir, href)
|
||
html = zf.read(path)
|
||
soup = etree.fromstring(html)
|
||
title_tag = soup.find(".//{http://www.w3.org/1999/xhtml}title")
|
||
title = _text(title_tag) if title_tag is not None else f"Section {i + 1}"
|
||
headings.append(Heading(title=title, level=1, href=href))
|
||
except Exception:
|
||
continue
|
||
return headings if headings else None
|
||
|
||
|
||
def _fallback_spine(zf: zipfile.ZipFile, opf: etree._Element, opf_dir: str) -> list[Heading]:
|
||
"""最后兜底：直接读所有 xhtml 的 h1-h3"""
|
||
return _try_spine(zf, opf, opf_dir) or []
|
||
|
||
|
||
def flatten_headings(headings: list[Heading], level_offset: int = 0) -> list[dict]:
|
||
"""将树形目录展平为列表，包含 heading_path"""
|
||
result = []
|
||
for h in headings:
|
||
result.append({
|
||
"title": h.title,
|
||
"level": h.level + level_offset,
|
||
"href": h.href,
|
||
})
|
||
result.extend(flatten_headings(h.children, level_offset))
|
||
return result
|
||
|
||
|
||
def build_heading_paths(headings: list[dict], order_idx: int) -> list[dict]:
|
||
"""为每个扁平条目构建完整 heading_path（如 '第3章 > 3.2 前向传播'）"""
|
||
stack = []
|
||
result = []
|
||
for h in headings:
|
||
while stack and stack[-1]["level"] >= h["level"]:
|
||
stack.pop()
|
||
stack.append(h)
|
||
path = " > ".join(s["title"] for s in stack)
|
||
result.append({**h, "heading_path": path, "order_idx": order_idx})
|
||
order_idx += 1
|
||
return result
|
||
|
||
|
||
def _resolve_path(opf_dir: str, href: str) -> str:
|
||
if href.startswith("/"):
|
||
return href.lstrip("/")
|
||
parts = [p for p in opf_dir.split("/") if p] if opf_dir else []
|
||
parts.append(href)
|
||
while ".." in parts:
|
||
idx = parts.index("..")
|
||
if idx > 0:
|
||
parts.pop(idx)
|
||
parts.pop(idx - 1)
|
||
else:
|
||
parts.pop(idx)
|
||
return "/".join(parts)
|
||
|
||
|
||
def _dirname(path: str) -> str:
|
||
parts = path.rsplit("/", 1)
|
||
return parts[0] if len(parts) > 1 else ""
|
||
|
||
|
||
def _text(el) -> str:
|
||
return "".join(el.itertext()).strip() if el is not None else ""
|
||
|
||
|
||
def href_to_fragment(href: str) -> str:
|
||
"""从 href 中提取 file 和 fragment 部分"""
|
||
if "#" in href:
|
||
return href.split("#", 1)[1]
|
||
return "" |