zwt13703
140 lines
4.3 KiB
Python
140 lines
4.3 KiB
Python
"""Word 文档解析引擎"""
|
|
from io import BytesIO
|
|
from typing import Any, Dict, Iterable, List
|
|
|
|
from docx import Document
|
|
from docx.document import Document as DocxDocument
|
|
from docx.oxml.table import CT_Tbl
|
|
from docx.oxml.text.paragraph import CT_P
|
|
from docx.table import Table as DocxTable
|
|
from docx.table import _Cell
|
|
from docx.text.paragraph import Paragraph
|
|
|
|
|
|
def parse_document(source: str | bytes) -> List[Dict[str, Any]]:
|
|
"""
|
|
解析 Word 文档，返回带 block_id、parent_block_id、sort_order 的 block 列表。
|
|
"""
|
|
doc = Document(BytesIO(source)) if isinstance(source, bytes) else Document(source)
|
|
blocks: List[Dict[str, Any]] = []
|
|
|
|
for item in iter_block_items(doc):
|
|
if isinstance(item, Paragraph):
|
|
block = classify_paragraph(item)
|
|
else:
|
|
block = classify_table(item)
|
|
|
|
if block:
|
|
blocks.append(block)
|
|
|
|
assign_block_ids(blocks)
|
|
build_tree(blocks)
|
|
return blocks
|
|
|
|
|
|
def iter_block_items(parent: DocxDocument | _Cell) -> Iterable[Paragraph | DocxTable]:
|
|
"""按 Word 文档中的实际顺序遍历段落和表格。"""
|
|
parent_element = parent.element.body if isinstance(parent, DocxDocument) else parent._tc
|
|
|
|
for child in parent_element.iterchildren():
|
|
if isinstance(child, CT_P):
|
|
yield Paragraph(child, parent)
|
|
elif isinstance(child, CT_Tbl):
|
|
yield DocxTable(child, parent)
|
|
|
|
|
|
def classify_paragraph(para: Paragraph) -> Dict[str, Any] | None:
|
|
"""识别段落类型"""
|
|
text = para.text.strip()
|
|
style_name = para.style.name if para.style else ""
|
|
|
|
if not text:
|
|
return None # 跳过空段落
|
|
|
|
if is_heading_style(style_name):
|
|
level = parse_heading_level(style_name)
|
|
return {"type": "heading", "text": text, "level": level, "style": style_name}
|
|
|
|
return {"type": "paragraph", "text": text, "level": 0, "style": style_name}
|
|
|
|
|
|
def is_heading_style(style_name: str) -> bool:
|
|
return style_name.startswith("Heading") or style_name.startswith("标题")
|
|
|
|
|
|
def parse_heading_level(style_name: str) -> int:
|
|
parts = style_name.replace("Heading", "").replace("标题", "").strip()
|
|
return int(parts) if parts.isdigit() and 1 <= int(parts) <= 6 else 1
|
|
|
|
|
|
def classify_table(table: DocxTable) -> Dict[str, Any] | None:
|
|
"""识别表格类型并提取行列数。"""
|
|
rows = len(table.rows)
|
|
cols = len(table.columns)
|
|
if rows == 0 or cols == 0:
|
|
return None
|
|
|
|
first_row = [cell.text.strip() for cell in table.rows[0].cells]
|
|
text = " | ".join(cell for cell in first_row if cell)
|
|
if not text:
|
|
text = f"表格（{rows} 行 x {cols} 列）"
|
|
|
|
return {
|
|
"type": "table",
|
|
"text": text,
|
|
"level": 0,
|
|
"rows": rows,
|
|
"cols": cols,
|
|
"headers": first_row,
|
|
}
|
|
|
|
|
|
def generate_block_id(index: int) -> str:
|
|
"""生成 block_id，格式：block_001"""
|
|
return f"block_{index:03d}"
|
|
|
|
|
|
def assign_block_ids(blocks: List[Dict[str, Any]]) -> None:
|
|
for index, block in enumerate(blocks, start=1):
|
|
block["block_id"] = generate_block_id(index)
|
|
block["sort_order"] = index
|
|
block["parent_block_id"] = None
|
|
|
|
|
|
def build_tree(blocks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
"""构建层级树结构"""
|
|
tree: List[Dict[str, Any]] = []
|
|
heading_stack: List[Dict[str, Any]] = []
|
|
|
|
for block in blocks:
|
|
node = {
|
|
"block_id": block["block_id"],
|
|
"block_type": block["type"],
|
|
"text": block["text"][:50],
|
|
"level": block.get("level", 0),
|
|
"children": [],
|
|
}
|
|
|
|
if block["type"] == "heading":
|
|
level = block.get("level", 1)
|
|
|
|
while heading_stack and heading_stack[-1]["level"] >= level:
|
|
heading_stack.pop()
|
|
|
|
if heading_stack:
|
|
parent = heading_stack[-1]["node"]
|
|
parent["children"].append(node)
|
|
block["parent_block_id"] = parent["block_id"]
|
|
else:
|
|
tree.append(node)
|
|
|
|
heading_stack.append({"level": level, "node": node})
|
|
elif heading_stack:
|
|
parent = heading_stack[-1]["node"]
|
|
parent["children"].append(node)
|
|
block["parent_block_id"] = parent["block_id"]
|
|
else:
|
|
tree.append(node)
|
|
|
|
return tree
|