Python脚本实现PDF书签的修改与导入导出

前言

在 使用PDFBookmark-Exchanger对PDF书签导入导出 中用了个Java写的软件导入导出。臃肿又不好看！还不如TerminalUI呢。遂拿Python重新写了一个。感谢Claude Opus 4.6。

代码

需要PyMuPDF才能正常运行。

#!/usr/bin/env python3
"""
PDF Bookmark Tool — 从 PDF 提取书签 / 从文件导入书签，预览、导出，并写入另一个 PDF。

依赖:
    pip install PyMuPDF rich

用法:
    python pdf_bookmark_tool.py
"""

import json
import re
import sys
import os

try:
    import fitz  # PyMuPDF
except ImportError:
    print("请先安装 PyMuPDF: pip install PyMuPDF")
    sys.exit(1)

try:
    from rich.console import Console
    from rich.tree import Tree
    from rich.prompt import Prompt, Confirm
    from rich.panel import Panel
    from rich.table import Table
    from rich import print as rprint
except ImportError:
    print("请先安装 rich: pip install rich")
    sys.exit(1)

console = Console()


# ──────────────────────────────────────────────
#  数据结构: 书签节点
# ──────────────────────────────────────────────
class BookmarkNode:
    """表示一个书签条目，支持嵌套子书签。"""

    def __init__(self, title: str, page: int, children: list["BookmarkNode"] | None = None):
        self.title = title
        self.page = page  # 0-based 页码
        self.children = children or []

    def to_dict(self) -> dict:
        return {
            "title": self.title,
            "page": self.page,
            "children": [c.to_dict() for c in self.children],
        }

    @staticmethod
    def from_dict(d: dict) -> "BookmarkNode":
        return BookmarkNode(
            title=d["title"],
            page=d["page"],
            children=[BookmarkNode.from_dict(c) for c in d.get("children", [])],
        )

    def to_text_lines(self, indent: int = 0) -> list[str]:
        """生成带缩进的纯文本行。"""
        prefix = "    " * indent
        lines = [f"{prefix}{self.title}  [p.{self.page + 1}]"]
        for child in self.children:
            lines.extend(child.to_text_lines(indent + 1))
        return lines


# ──────────────────────────────────────────────
#  从 PDF 提取书签
# ──────────────────────────────────────────────
def extract_bookmarks(pdf_path: str) -> list[BookmarkNode]:
    """从 PDF 文件提取书签树（TOC）。"""
    doc = fitz.open(pdf_path)
    toc = doc.get_toc(simple=True)
    doc.close()

    if not toc:
        return []

    root: list[BookmarkNode] = []
    stack: list[tuple[int, BookmarkNode]] = []

    for level, title, page in toc:
        node = BookmarkNode(title=title, page=page - 1)

        while stack and stack[-1][0] >= level:
            stack.pop()

        if stack:
            stack[-1][1].children.append(node)
        else:
            root.append(node)

        stack.append((level, node))

    return root


# ──────────────────────────────────────────────
#  从 TXT 导入书签
# ──────────────────────────────────────────────
def import_from_txt(txt_path: str) -> list[BookmarkNode]:
    """
    从带缩进的 TXT 文件导入书签。

    每行格式:
        <缩进><标题>  [p.<页码>]

    缩进使用 4 个空格或 1 个制表符表示一级。
    """
    pattern = re.compile(
        r"^(?P<indent>[\t ]*)(?P<title>.+?)\s+\[p\.(?P<page>\d+)\]\s*$"
    )

    nodes_with_level: list[tuple[int, BookmarkNode]] = []

    with open(txt_path, "r", encoding="utf-8") as f:
        for line_no, raw_line in enumerate(f, start=1):
            line = raw_line.rstrip("\n\r")
            if not line.strip():
                continue

            m = pattern.match(line)
            if not m:
                console.print(
                    f"  [yellow]⚠️  第 {line_no} 行格式不匹配，已跳过: "
                    f"{line.strip()!r}[/yellow]"
                )
                continue

            indent_str = m.group("indent")
            if "\t" in indent_str:
                level = indent_str.count("\t") + 1
            else:
                level = len(indent_str) // 4 + 1

            title = m.group("title").strip()
            page = int(m.group("page")) - 1

            nodes_with_level.append((level, BookmarkNode(title=title, page=page)))

    return _build_tree_from_flat(nodes_with_level)


def _build_tree_from_flat(items: list[tuple[int, BookmarkNode]]) -> list[BookmarkNode]:
    """将 (level, node) 扁平列表构建为嵌套树。"""
    root: list[BookmarkNode] = []
    stack: list[tuple[int, BookmarkNode]] = []

    for level, node in items:
        while stack and stack[-1][0] >= level:
            stack.pop()

        if stack:
            stack[-1][1].children.append(node)
        else:
            root.append(node)

        stack.append((level, node))

    return root


# ──────────────────────────────────────────────
#  从 JSON 导入书签
# ──────────────────────────────────────────────
def import_from_json(json_path: str) -> list[BookmarkNode]:
    """从 JSON 文件导入书签。"""
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    if not isinstance(data, list):
        raise ValueError("JSON 根元素应为数组")

    return [BookmarkNode.from_dict(d) for d in data]


# ──────────────────────────────────────────────
#  写入书签到 PDF
# ──────────────────────────────────────────────
def _flatten_bookmarks(nodes: list[BookmarkNode], level: int = 1) -> list[list]:
    result = []
    for node in nodes:
        result.append([level, node.title, node.page + 1])
        result.extend(_flatten_bookmarks(node.children, level + 1))
    return result


def write_bookmarks(target_pdf_path: str, bookmarks: list[BookmarkNode], output_path: str):
    doc = fitz.open(target_pdf_path)
    toc = _flatten_bookmarks(bookmarks)
    doc.set_toc(toc)
    if os.path.abspath(target_pdf_path) == os.path.abspath(output_path):
        doc.save(output_path, incremental=True, encryption=0)
    else:
        doc.save(output_path)
    doc.close()


# ──────────────────────────────────────────────
#  导出
# ──────────────────────────────────────────────
def export_as_txt(bookmarks: list[BookmarkNode], output_path: str):
    lines = []
    for node in bookmarks:
        lines.extend(node.to_text_lines())
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")


def export_as_json(bookmarks: list[BookmarkNode], output_path: str):
    data = [node.to_dict() for node in bookmarks]
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


# ──────────────────────────────────────────────
#  TUI 预览（支持行数限制）
# ──────────────────────────────────────────────
def _collect_flat_lines(nodes: list[BookmarkNode], indent: int = 0) -> list[tuple[int, BookmarkNode]]:
    """将嵌套书签树按深度优先展开为 (indent_level, node) 的扁平列表。"""
    result = []
    for node in nodes:
        result.append((indent, node))
        result.extend(_collect_flat_lines(node.children, indent + 1))
    return result


def preview_bookmarks(
    bookmarks: list[BookmarkNode],
    title: str = "📑 书签预览",
    max_lines: int | None = None,
):
    """
    在终端中以树状结构预览书签。

    Args:
        bookmarks: 书签列表
        title:     树标题
        max_lines: 最多显示行数，None 表示全部显示
    """
    flat = _collect_flat_lines(bookmarks)
    total_lines = len(flat)
    truncated = False

    if max_lines is not None and total_lines > max_lines:
        display = flat[:max_lines]
        truncated = True
    else:
        display = flat

    # 构建只包含 display 中节点的树
    tree = Tree(f"[bold cyan]{title}[/bold cyan]")
    # 用栈追踪每一层当前的 Tree 分支
    branch_stack: list[tuple[int, Tree]] = [(-1, tree)]  # (indent_level, tree_node)

    for indent, node in display:
        label = f"[bold]{node.title}[/bold]  [dim](p.{node.page + 1})[/dim]"

        # 回退到正确的父级
        while branch_stack and branch_stack[-1][0] >= indent:
            branch_stack.pop()

        parent_tree = branch_stack[-1][1] if branch_stack else tree
        branch = parent_tree.add(label)
        branch_stack.append((indent, branch))

    console.print(tree)

    if truncated:
        remaining = total_lines - max_lines
        console.print(
            f"[dim]... 还有 {remaining} 个条目未显示，"
            f"选择「重新预览书签」可查看全部 {total_lines} 条 ...[/dim]"
        )


def count_bookmarks(nodes: list[BookmarkNode]) -> int:
    total = len(nodes)
    for node in nodes:
        total += count_bookmarks(node.children)
    return total


# ──────────────────────────────────────────────
#  来源选择：从 PDF / TXT / JSON 加载书签
# ──────────────────────────────────────────────
def load_bookmarks_interactive() -> tuple[list[BookmarkNode], str]:
    table = Table(title="选择书签来源", show_header=False, border_style="blue")
    table.add_column("选项", style="bold yellow", width=1)
    table.add_column("描述")
    table.add_row("1", "从 PDF  文件提取书签")
    table.add_row("2", "从 TXT  文件导入书签")
    table.add_row("3", "从 JSON 文件导入书签")
    console.print(table)

    choice = Prompt.ask("[bold]请选择来源[/bold]", choices=["1", "2", "3"])

    if choice == "1":
        path = Prompt.ask("[bold green]📂 请输入含有书签的源 PDF 路径[/bold green]").strip().strip('"').strip("'")
        if not os.path.isfile(path):
            console.print(f"[bold red]❌ 文件不存在: {path}[/bold red]")
            sys.exit(1)
        console.print("[cyan]正在提取书签...[/cyan]")
        bookmarks = extract_bookmarks(path)
        source_desc = os.path.basename(path)

    elif choice == "2":
        path = Prompt.ask("[bold green]📂 请输入 TXT 书签文件路径[/bold green]").strip().strip('"').strip("'")
        if not os.path.isfile(path):
            console.print(f"[bold red]❌ 文件不存在: {path}[/bold red]")
            sys.exit(1)
        console.print("[cyan]正在导入 TXT 书签...[/cyan]")
        bookmarks = import_from_txt(path)
        source_desc = os.path.basename(path)

    elif choice == "3":
        path = Prompt.ask("[bold green]📂 请输入 JSON 书签文件路径[/bold green]").strip().strip('"').strip("'")
        if not os.path.isfile(path):
            console.print(f"[bold red]❌ 文件不存在: {path}[/bold red]")
            sys.exit(1)
        console.print("[cyan]正在导入 JSON 书签...[/cyan]")
        try:
            bookmarks = import_from_json(path)
        except (json.JSONDecodeError, ValueError) as e:
            console.print(f"[bold red]❌ JSON 解析失败: {e}[/bold red]")
            sys.exit(1)
        source_desc = os.path.basename(path)

    if not bookmarks:
        console.print("[bold red]⚠️  未找到任何书签条目。[/bold red]")
        sys.exit(1)

    return bookmarks, source_desc


# ──────────────────────────────────────────────
#  主交互循环 (TUI)
# ──────────────────────────────────────────────
INITIAL_PREVIEW_LINES = 16


def main():
    console.print(
        Panel.fit(
            "[bold magenta]PDF Bookmark Tool[/bold magenta]\n"
            "[dim]提取 · 导入 · 预览 · 导出 · 写入 PDF 书签[/dim]",
            border_style="bright_blue",
        )
    )

    # ── 步骤 1: 选择来源并加载书签 ──
    bookmarks, source_desc = load_bookmarks_interactive()

    total = count_bookmarks(bookmarks)
    console.print(f"[bold green]✅ 成功加载 {total} 个书签条目[/bold green]\n")

    # ── 步骤 2: 首次预览（仅显示前 16 行）──
    preview_bookmarks(
        bookmarks,
        title=f"书签预览 — {source_desc}",
        max_lines=INITIAL_PREVIEW_LINES,
    )

    # ── 步骤 3: 操作菜单 ──
    while True:
        console.print()
        table = Table(title="操作菜单", show_header=False, border_style="blue")
        table.add_column("选项", style="bold yellow", width=6)
        table.add_column("描述")
        table.add_row("1", "导出书签为 TXT（带缩进）")
        table.add_row("2", "导出书签为 JSON")
        table.add_row("3", "将书签写入目标 PDF")
        table.add_row("4", "重新预览书签（完整）")
        table.add_row("5", "重新加载 / 切换书签来源")
        table.add_row("q", "退出")
        console.print(table)

        choice = Prompt.ask(
            "[bold]请选择操作[/bold]",
            choices=["1", "2", "3", "4", "5", "q"],
            default="q",
        )

        if choice == "1":
            out = Prompt.ask("  导出 TXT 路径", default="bookmarks.txt").strip().strip('"').strip("'")
            export_as_txt(bookmarks, out)
            console.print(f"  [green]✅ 已导出到 {out}[/green]")

        elif choice == "2":
            out = Prompt.ask("  导出 JSON 路径", default="bookmarks.json").strip().strip('"').strip("'")
            export_as_json(bookmarks, out)
            console.print(f"  [green]✅ 已导出到 {out}[/green]")

        elif choice == "3":
            target = Prompt.ask("  📂 请输入目标 PDF 路径").strip().strip('"').strip("'")
            if not os.path.isfile(target):
                console.print(f"  [red]❌ 文件不存在: {target}[/red]")
                continue
            save_path = (
                Prompt.ask("  💾 保存路径（直接回车则覆盖原文件）", default=target)
                .strip()
                .strip('"')
                .strip("'")
            )

            if os.path.abspath(save_path) == os.path.abspath(target):
                overwrite = Confirm.ask(
                    f"  🤔  将覆盖 [bold]{target}[/bold]，确认？", default=False
                )
                if not overwrite:
                    continue

            try:
                write_bookmarks(target, bookmarks, save_path)
                console.print(f"  [green]✅ 书签已写入 {save_path}[/green]")
            except Exception as e:
                console.print(f"  [red]❌ 写入失败: {e}[/red]")

        elif choice == "4":
            # 完整预览，不截断
            preview_bookmarks(bookmarks, title=f"书签预览（完整） — {source_desc}")

        elif choice == "5":
            bookmarks, source_desc = load_bookmarks_interactive()
            total = count_bookmarks(bookmarks)
            console.print(f"[bold green]✅ 成功加载 {total} 个书签条目[/bold green]\n")
            preview_bookmarks(
                bookmarks,
                title=f"书签预览 — {source_desc}",
                max_lines=INITIAL_PREVIEW_LINES,
            )

        elif choice == "q":
            console.print("[dim]再见！🥰[/dim]")
            break


if __name__ == "__main__":
    main()

用起来还不错！就是打包成exe后体积大了不少，比原先用Java实现的多了亿点空间。