337 lines
11 KiB
Python
337 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
PDF目录摘取测试脚本 (与Java代码逻辑一致)
|
||
|
||
模拟 FileServiceImpl.extractDirectoryFromPDF 的逻辑，用于测试验证。
|
||
|
||
使用方法:
|
||
python test_catalog_extract.py <pdf文件路径>
|
||
|
||
依赖安装:
|
||
pip install pypdf
|
||
"""
|
||
|
||
import sys
|
||
import json
|
||
import os
|
||
from datetime import datetime
|
||
|
||
# 模拟Java的日志输出
|
||
def log_info(msg):
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] INFO [目录摘取] {msg}")
|
||
|
||
def log_debug(msg):
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] DEBUG [目录摘取] {msg}")
|
||
|
||
def log_warn(msg):
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] WARN [目录摘取] {msg}")
|
||
|
||
def log_error(msg):
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] ERROR [目录摘取] {msg}")
|
||
|
||
|
||
def get_page_number_from_outline_item(reader, item):
|
||
"""
|
||
对应Java: getPageNumberFromOutlineItem(PDOutlineItem item)
|
||
尝试从目录项获取页码
|
||
"""
|
||
try:
|
||
if hasattr(item, 'page') and item.page:
|
||
# 直接有page属性
|
||
page_num = reader.get_destination_page_number(item)
|
||
return page_num + 1 # 转为1-based
|
||
except:
|
||
pass
|
||
return -1
|
||
|
||
|
||
def extract_outline_items(reader, outline_items, chapters, level=1, indent=""):
|
||
"""
|
||
对应Java: extractOutlineItems(PDOutlineItem item, String indent, List<JSONObject> chapters, int level)
|
||
递归提取PDF目录项
|
||
"""
|
||
if not outline_items:
|
||
return
|
||
|
||
for i, item in enumerate(outline_items):
|
||
try:
|
||
# 如果是列表（子项），递归处理
|
||
if isinstance(item, list):
|
||
extract_outline_items(reader, item, chapters, level + 1, indent + " ")
|
||
continue
|
||
|
||
# 获取标题
|
||
title = item.title if hasattr(item, 'title') else None
|
||
|
||
if title and title.strip():
|
||
title = title.strip()
|
||
|
||
chapter = {
|
||
"chapterId": f"chap-{len(chapters)}",
|
||
"title": title,
|
||
"level": level,
|
||
"page": "1",
|
||
"startPage": 1,
|
||
"endPage": 1,
|
||
"children": []
|
||
}
|
||
|
||
# 尝试获取页码信息
|
||
start_page = -1
|
||
try:
|
||
page_num = reader.get_destination_page_number(item)
|
||
if page_num is not None and page_num >= 0:
|
||
start_page = page_num + 1 # PyPDF页码从0开始，转为1-based
|
||
except Exception as e:
|
||
log_debug(f"无法获取章节页码信息: {title} - {e}")
|
||
|
||
# 如果获取到了起始页码
|
||
if start_page != -1:
|
||
chapter["page"] = str(start_page)
|
||
chapter["startPage"] = start_page
|
||
chapter["endPage"] = start_page # 暂时设为起始页码
|
||
else:
|
||
chapter["page"] = "1"
|
||
chapter["startPage"] = 1
|
||
chapter["endPage"] = 1
|
||
|
||
chapters.append(chapter)
|
||
log_debug(f"{indent}• {title} (第{chapter['startPage']}页)")
|
||
|
||
except Exception as e:
|
||
log_warn(f"处理目录项失败: {e}")
|
||
|
||
|
||
def build_hierarchical_structure(chapters, max_depth=2):
|
||
"""
|
||
对应Java: buildHierarchicalStructure(List<JSONObject> chapters)
|
||
构建层级目录结构
|
||
"""
|
||
root_array = []
|
||
if not chapters:
|
||
return root_array
|
||
|
||
level_parents = [] # 存储每个层级的当前父节点
|
||
|
||
for chapter in chapters:
|
||
level = max(0, chapter.get("level", 1))
|
||
if chapter.get("children") is None:
|
||
chapter["children"] = []
|
||
|
||
if level == 0 or not level_parents or level_parents[0] is None:
|
||
root_array.append(chapter)
|
||
level_parents.clear()
|
||
level_parents.append(chapter)
|
||
continue
|
||
|
||
parent_index = min(level - 1, len(level_parents) - 1)
|
||
parent = level_parents[parent_index] if parent_index >= 0 else None
|
||
|
||
if parent is None:
|
||
# 找不到父节点时降级为根节点
|
||
root_array.append(chapter)
|
||
level_parents.clear()
|
||
level_parents.append(chapter)
|
||
continue
|
||
|
||
parent_children = parent.get("children", [])
|
||
parent_children.append(chapter)
|
||
parent["children"] = parent_children
|
||
|
||
# 更新level_parents
|
||
while len(level_parents) < level:
|
||
level_parents.append(None)
|
||
if level >= len(level_parents):
|
||
level_parents.append(chapter)
|
||
else:
|
||
level_parents[level] = chapter
|
||
|
||
# 清空更深层级的父节点
|
||
for idx in range(level + 1, len(level_parents)):
|
||
level_parents[idx] = None
|
||
|
||
# 限制目录深度
|
||
limit_directory_depth(root_array, max_depth)
|
||
|
||
return root_array
|
||
|
||
|
||
def limit_directory_depth(chapters, max_depth, current_depth=1):
|
||
"""限制目录深度"""
|
||
if current_depth >= max_depth:
|
||
for chapter in chapters:
|
||
chapter["children"] = []
|
||
else:
|
||
for chapter in chapters:
|
||
if chapter.get("children"):
|
||
limit_directory_depth(chapter["children"], max_depth, current_depth + 1)
|
||
|
||
|
||
def extract_directory_from_pdf(file_path, file_id="test-file", file_name=None):
|
||
"""
|
||
对应Java: extractDirectoryFromPDF(String filePath, String fileId, String fileName)
|
||
从PDF中提取目录结构
|
||
"""
|
||
try:
|
||
from pypdf import PdfReader
|
||
except ImportError:
|
||
try:
|
||
from PyPDF2 import PdfReader
|
||
except ImportError:
|
||
log_error("请先安装 pypdf: pip install pypdf")
|
||
return None
|
||
|
||
if file_name is None:
|
||
file_name = os.path.basename(file_path)
|
||
|
||
log_info(f"开始处理PDF文件: {file_name}")
|
||
|
||
try:
|
||
total_pages = 1
|
||
chapters = []
|
||
|
||
# 加载PDF文档
|
||
reader = PdfReader(file_path)
|
||
total_pages = len(reader.pages)
|
||
log_info(f"PDF总页数: {total_pages}")
|
||
|
||
# 尝试从PDF大纲中提取目录
|
||
outline = reader.outline
|
||
|
||
if outline:
|
||
log_info("发现PDF大纲，开始提取...")
|
||
extract_outline_items(reader, outline, chapters, level=1, indent="")
|
||
log_info(f"从大纲提取到 {len(chapters)} 个章节")
|
||
|
||
if chapters:
|
||
# 构建层级结构
|
||
hierarchical_array = build_hierarchical_structure(chapters)
|
||
log_info(f"成功构建层级目录结构，根节点数: {len(hierarchical_array)}")
|
||
return {
|
||
"source": "pdf_outline",
|
||
"total_pages": total_pages,
|
||
"flat_chapters": chapters,
|
||
"hierarchical_chapters": hierarchical_array
|
||
}
|
||
|
||
# 回退方案：使用文件名作为单个章节
|
||
title = file_name
|
||
if title and title.lower().endswith(".pdf"):
|
||
title = title[:-4]
|
||
if not title or not title.strip():
|
||
title = "文档内容"
|
||
|
||
directory_array = [{
|
||
"chapterId": "chap-0",
|
||
"title": title,
|
||
"level": 1,
|
||
"page": "1",
|
||
"startPage": 1,
|
||
"endPage": total_pages,
|
||
"children": []
|
||
}]
|
||
|
||
log_info(f"使用文件名作为默认章节: '{title}', 页数范围: 1-{total_pages}")
|
||
|
||
return {
|
||
"source": "filename_fallback",
|
||
"total_pages": total_pages,
|
||
"flat_chapters": directory_array,
|
||
"hierarchical_chapters": directory_array
|
||
}
|
||
|
||
except Exception as e:
|
||
log_error(f"处理PDF失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return None
|
||
|
||
|
||
def print_catalog_tree(chapters, indent=0):
|
||
"""打印目录树形结构"""
|
||
for chapter in chapters:
|
||
prefix = " " * indent + ("├─ " if indent > 0 else "")
|
||
title = chapter.get("title", "未命名")
|
||
page = chapter.get("startPage", "?")
|
||
level = chapter.get("level", 1)
|
||
print(f"{prefix}[L{level}] {title} (第{page}页)")
|
||
|
||
children = chapter.get("children", [])
|
||
if children:
|
||
print_catalog_tree(children, indent + 1)
|
||
|
||
|
||
def main():
|
||
print("=" * 70)
|
||
print("PDF目录摘取测试脚本 (与Java extractDirectoryFromPDF 逻辑一致)")
|
||
print("=" * 70)
|
||
|
||
if len(sys.argv) < 2:
|
||
print("\n用法:")
|
||
print(" python test_catalog_extract.py <pdf文件路径>")
|
||
print("\n示例:")
|
||
print(" python test_catalog_extract.py ./document.pdf")
|
||
print(" python test_catalog_extract.py /path/to/技术方案.pdf")
|
||
print("\n依赖:")
|
||
print(" pip install pypdf")
|
||
print("=" * 70)
|
||
sys.exit(1)
|
||
|
||
pdf_path = sys.argv[1]
|
||
|
||
# 检查文件是否存在
|
||
if not os.path.exists(pdf_path):
|
||
log_error(f"文件不存在: {pdf_path}")
|
||
sys.exit(1)
|
||
|
||
print(f"\n📄 PDF文件: {pdf_path}")
|
||
print(f"⏰ 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print("-" * 70)
|
||
|
||
start_time = datetime.now()
|
||
|
||
# 提取目录
|
||
result = extract_directory_from_pdf(pdf_path)
|
||
|
||
duration = (datetime.now() - start_time).total_seconds()
|
||
|
||
if result is None:
|
||
log_error("目录提取失败")
|
||
sys.exit(1)
|
||
|
||
print("-" * 70)
|
||
print(f"\n📊 提取结果:")
|
||
print(f" 来源: {result.get('source', 'unknown')}")
|
||
print(f" 总页数: {result.get('total_pages', '?')}")
|
||
print(f" 扁平目录项数量: {len(result.get('flat_chapters', []))}")
|
||
print(f" 层级根节点数量: {len(result.get('hierarchical_chapters', []))}")
|
||
print(f" 处理耗时: {duration:.3f}秒")
|
||
|
||
# 打印扁平目录
|
||
print(f"\n📋 扁平目录 (flat_chapters):")
|
||
print("-" * 40)
|
||
for chapter in result.get("flat_chapters", []):
|
||
level = chapter.get("level", 1)
|
||
indent = " " * (level - 1)
|
||
print(f"{indent}[{chapter.get('chapterId')}] {chapter.get('title')} (L{level}, 第{chapter.get('startPage')}页)")
|
||
|
||
# 打印层级目录
|
||
print(f"\n🌳 层级目录 (hierarchical_chapters):")
|
||
print("-" * 40)
|
||
print_catalog_tree(result.get("hierarchical_chapters", []))
|
||
|
||
# 保存JSON结果
|
||
output_file = os.path.splitext(pdf_path)[0] + "_catalog_test.json"
|
||
with open(output_file, 'w', encoding='utf-8') as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||
print(f"\n✅ JSON结果已保存到: {output_file}")
|
||
|
||
print("\n" + "=" * 70)
|
||
print(f"⏰ 完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print("=" * 70)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|