478 lines
16 KiB
Python
478 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import re
|
||
import tempfile
|
||
import zipfile
|
||
from pathlib import Path
|
||
from typing import Any, Iterable
|
||
|
||
from _document_builder import (
|
||
add_blocks,
|
||
apply_core_properties,
|
||
apply_header_footer,
|
||
apply_page_settings,
|
||
expect_list,
|
||
expect_object,
|
||
set_update_fields,
|
||
)
|
||
from _docx_common import (
|
||
DOCX_INPUT_SUFFIXES,
|
||
SkillArgumentParser,
|
||
W_NS,
|
||
input_file,
|
||
inspect_archive,
|
||
load_json_argument,
|
||
output_file,
|
||
parse_xml_bytes,
|
||
publish_file,
|
||
run_cli,
|
||
)
|
||
|
||
|
||
MAX_OPERATIONS = 200
|
||
MAX_REPLACEMENTS = 10_000
|
||
|
||
|
||
def _has_revisions(source: Path) -> bool:
|
||
revision_tags = {
|
||
f"{{{W_NS}}}{name}"
|
||
for name in (
|
||
"ins",
|
||
"del",
|
||
"moveFrom",
|
||
"moveTo",
|
||
"pPrChange",
|
||
"rPrChange",
|
||
"tblPrChange",
|
||
"trPrChange",
|
||
"tcPrChange",
|
||
"sectPrChange",
|
||
)
|
||
}
|
||
with zipfile.ZipFile(source) as archive:
|
||
for name in archive.namelist():
|
||
if not name.startswith("word/") or not name.endswith(".xml"):
|
||
continue
|
||
root = parse_xml_bytes(archive.read(name), label=name)
|
||
if any(element.tag in revision_tags for element in root.iter()):
|
||
return True
|
||
return False
|
||
|
||
|
||
def _iter_table_paragraphs(table: Any) -> Iterable[Any]:
|
||
for row in table.rows:
|
||
for cell in row.cells:
|
||
yield from cell.paragraphs
|
||
for nested_table in cell.tables:
|
||
yield from _iter_table_paragraphs(nested_table)
|
||
|
||
|
||
def _iter_scope_paragraphs(document: Any, scope: str) -> list[Any]:
|
||
if scope not in {"body", "tables", "headers", "footers", "all"}:
|
||
raise ValueError("scope 仅支持 body、tables、headers、footers、all")
|
||
result: list[Any] = []
|
||
if scope in {"body", "all"}:
|
||
result.extend(document.paragraphs)
|
||
if scope in {"tables", "all"}:
|
||
for table in document.tables:
|
||
result.extend(_iter_table_paragraphs(table))
|
||
if scope in {"headers", "all"}:
|
||
for section in document.sections:
|
||
result.extend(section.header.paragraphs)
|
||
for table in section.header.tables:
|
||
result.extend(_iter_table_paragraphs(table))
|
||
if scope in {"footers", "all"}:
|
||
for section in document.sections:
|
||
result.extend(section.footer.paragraphs)
|
||
for table in section.footer.tables:
|
||
result.extend(_iter_table_paragraphs(table))
|
||
deduplicated: list[Any] = []
|
||
seen: set[int] = set()
|
||
for paragraph in result:
|
||
identity = id(paragraph._p)
|
||
if identity not in seen:
|
||
seen.add(identity)
|
||
deduplicated.append(paragraph)
|
||
return deduplicated
|
||
|
||
|
||
def _matches(
|
||
text: str,
|
||
find: str,
|
||
*,
|
||
match_case: bool,
|
||
whole_word: bool,
|
||
) -> list[re.Match[str]]:
|
||
pattern = re.escape(find)
|
||
if whole_word:
|
||
pattern = rf"(?<!\w){pattern}(?!\w)"
|
||
flags = 0 if match_case else re.IGNORECASE
|
||
return list(re.finditer(pattern, text, flags))
|
||
|
||
|
||
def _replace_in_paragraph(
|
||
paragraph: Any,
|
||
*,
|
||
find: str,
|
||
replacement: str,
|
||
match_case: bool,
|
||
whole_word: bool,
|
||
remaining: int,
|
||
tracker: Any = None,
|
||
) -> int:
|
||
runs = paragraph.runs
|
||
if not runs:
|
||
return 0
|
||
texts = [run.text for run in runs]
|
||
full_text = "".join(texts)
|
||
found = _matches(
|
||
full_text,
|
||
find,
|
||
match_case=match_case,
|
||
whole_word=whole_word,
|
||
)
|
||
if remaining >= 0:
|
||
found = found[:remaining]
|
||
if not found:
|
||
return 0
|
||
if tracker is not None:
|
||
tracker.replace(paragraph, found, replacement)
|
||
return len(found)
|
||
|
||
spans: list[tuple[int, int]] = []
|
||
offset = 0
|
||
for text in texts:
|
||
spans.append((offset, offset + len(text)))
|
||
offset += len(text)
|
||
|
||
for match in reversed(found):
|
||
start, end = match.span()
|
||
start_index = next(
|
||
index
|
||
for index, (_, run_end) in enumerate(spans)
|
||
if start < run_end or (start == 0 and run_end == 0)
|
||
)
|
||
end_index = next(
|
||
index
|
||
for index, (run_start, run_end) in enumerate(spans)
|
||
if end <= run_end and end > run_start
|
||
)
|
||
start_run_start, _ = spans[start_index]
|
||
end_run_start, _ = spans[end_index]
|
||
start_offset = start - start_run_start
|
||
end_offset = end - end_run_start
|
||
if start_index == end_index:
|
||
current = runs[start_index].text
|
||
runs[start_index].text = (
|
||
current[:start_offset] + replacement + current[end_offset:]
|
||
)
|
||
else:
|
||
start_text = runs[start_index].text
|
||
end_text = runs[end_index].text
|
||
runs[start_index].text = start_text[:start_offset] + replacement
|
||
for index in range(start_index + 1, end_index):
|
||
runs[index].text = ""
|
||
runs[end_index].text = end_text[end_offset:]
|
||
return len(found)
|
||
|
||
|
||
def _op_replace_text(document: Any, op: dict[str, Any], tracker: Any = None) -> dict[str, Any]:
|
||
find = str(op.get("find", ""))
|
||
if not find:
|
||
raise ValueError("replace_text.find 不能为空")
|
||
replacement = str(op.get("replace", ""))
|
||
limit = int(op.get("count", 0))
|
||
if limit < 0 or limit > MAX_REPLACEMENTS:
|
||
raise ValueError(f"replace_text.count 必须在 0 到 {MAX_REPLACEMENTS} 之间")
|
||
scope = str(op.get("scope", "all"))
|
||
paragraphs = _iter_scope_paragraphs(document, scope)
|
||
total = 0
|
||
affected: list[int] = []
|
||
for index, paragraph in enumerate(paragraphs):
|
||
remaining = -1 if limit == 0 else limit - total
|
||
if remaining == 0:
|
||
break
|
||
count = _replace_in_paragraph(
|
||
paragraph,
|
||
find=find,
|
||
replacement=replacement,
|
||
match_case=bool(op.get("match_case", True)),
|
||
whole_word=bool(op.get("whole_word", False)),
|
||
remaining=remaining,
|
||
tracker=tracker,
|
||
)
|
||
if count:
|
||
total += count
|
||
if len(affected) < 100:
|
||
affected.append(index)
|
||
if bool(op.get("required", True)) and total == 0:
|
||
raise ValueError(f"未找到要替换的文本：{find!r}")
|
||
return {
|
||
"type": "replace_text",
|
||
"replacement_count": total,
|
||
"affected_paragraph_indexes": affected,
|
||
}
|
||
|
||
|
||
def _op_append_blocks(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
counts = add_blocks(document, document, op.get("blocks"))
|
||
return {"type": "append_blocks", "counts": counts}
|
||
|
||
|
||
def _find_body_paragraph(document: Any, text: str, match: str) -> Any:
|
||
for paragraph in document.paragraphs:
|
||
if match == "exact" and paragraph.text == text:
|
||
return paragraph
|
||
if match == "contains" and text in paragraph.text:
|
||
return paragraph
|
||
raise ValueError(f"未找到锚点段落：{text!r}")
|
||
|
||
|
||
def _op_insert_blocks_after(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
anchor_text = str(op.get("find", ""))
|
||
if not anchor_text:
|
||
raise ValueError("insert_blocks_after.find 不能为空")
|
||
match = str(op.get("match", "contains"))
|
||
if match not in {"exact", "contains"}:
|
||
raise ValueError("insert_blocks_after.match 仅支持 exact、contains")
|
||
anchor = _find_body_paragraph(document, anchor_text, match)
|
||
body = document.element.body
|
||
before_elements = set(body.iterchildren())
|
||
counts = add_blocks(document, document, op.get("blocks"))
|
||
new_elements = [
|
||
child
|
||
for child in body.iterchildren()
|
||
if child not in before_elements and child.tag != f"{{{W_NS}}}sectPr"
|
||
]
|
||
insertion_point = anchor._p
|
||
for element in new_elements:
|
||
body.remove(element)
|
||
insertion_point.addnext(element)
|
||
insertion_point = element
|
||
return {"type": "insert_blocks_after", "counts": counts}
|
||
|
||
|
||
def _op_remove_paragraphs(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
text = str(op.get("text", ""))
|
||
if not text:
|
||
raise ValueError("remove_paragraphs.text 不能为空")
|
||
match = str(op.get("match", "contains"))
|
||
if match not in {"exact", "contains"}:
|
||
raise ValueError("remove_paragraphs.match 仅支持 exact、contains")
|
||
limit = int(op.get("count", 0))
|
||
removed = 0
|
||
for paragraph in list(document.paragraphs):
|
||
matched = paragraph.text == text if match == "exact" else text in paragraph.text
|
||
if not matched:
|
||
continue
|
||
paragraph._element.getparent().remove(paragraph._element)
|
||
removed += 1
|
||
if limit and removed >= limit:
|
||
break
|
||
if bool(op.get("required", True)) and removed == 0:
|
||
raise ValueError(f"未找到要删除的段落：{text!r}")
|
||
return {"type": "remove_paragraphs", "removed_count": removed}
|
||
|
||
|
||
def _op_set_paragraph_style(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
style = str(op.get("style", ""))
|
||
if not style:
|
||
raise ValueError("set_paragraph_style.style 不能为空")
|
||
if style not in document.styles:
|
||
raise ValueError(f"文档样式不存在：{style}")
|
||
indexes = op.get("indexes")
|
||
changed = 0
|
||
if indexes is not None:
|
||
for raw_index in expect_list(indexes, "set_paragraph_style.indexes"):
|
||
index = int(raw_index)
|
||
if index < 0 or index >= len(document.paragraphs):
|
||
raise ValueError(f"段落索引超出范围：{index}")
|
||
document.paragraphs[index].style = style
|
||
changed += 1
|
||
else:
|
||
text = str(op.get("contains", ""))
|
||
if not text:
|
||
raise ValueError("需要提供 indexes 或 contains")
|
||
for paragraph in document.paragraphs:
|
||
if text in paragraph.text:
|
||
paragraph.style = style
|
||
changed += 1
|
||
if bool(op.get("required", True)) and changed == 0:
|
||
raise ValueError("没有段落应用到指定样式")
|
||
return {"type": "set_paragraph_style", "changed_count": changed}
|
||
|
||
|
||
def _op_set_properties(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
apply_core_properties(document, op.get("properties"))
|
||
return {"type": "set_properties"}
|
||
|
||
|
||
def _op_set_page(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
selection = op.get("section", "all")
|
||
if selection == "all":
|
||
sections = list(document.sections)
|
||
else:
|
||
index = int(selection)
|
||
if index < 0 or index >= len(document.sections):
|
||
raise ValueError(f"section 索引超出范围：{index}")
|
||
sections = [document.sections[index]]
|
||
for section in sections:
|
||
apply_page_settings(section, op.get("page"))
|
||
return {"type": "set_page", "section_count": len(sections)}
|
||
|
||
|
||
def _op_set_header_footer(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
apply_header_footer(
|
||
document,
|
||
op.get("header"),
|
||
op.get("footer"),
|
||
)
|
||
return {"type": "set_header_footer", "section_count": len(document.sections)}
|
||
|
||
|
||
def _op_remove_tables(document: Any, op: dict[str, Any]) -> dict[str, Any]:
|
||
indexes = sorted(
|
||
{int(item) for item in expect_list(op.get("indexes"), "remove_tables.indexes")},
|
||
reverse=True,
|
||
)
|
||
for index in indexes:
|
||
if index < 0 or index >= len(document.tables):
|
||
raise ValueError(f"表格索引超出范围：{index}")
|
||
table = document.tables[index]
|
||
table._element.getparent().remove(table._element)
|
||
return {"type": "remove_tables", "removed_count": len(indexes)}
|
||
|
||
|
||
def _apply_operation(document: Any, raw_op: Any) -> dict[str, Any]:
|
||
op = expect_object(raw_op, "operations[]")
|
||
kind = str(op.get("type", ""))
|
||
handlers = {
|
||
"replace_text": _op_replace_text,
|
||
"append_blocks": _op_append_blocks,
|
||
"insert_blocks_after": _op_insert_blocks_after,
|
||
"remove_paragraphs": _op_remove_paragraphs,
|
||
"set_paragraph_style": _op_set_paragraph_style,
|
||
"set_properties": _op_set_properties,
|
||
"set_page": _op_set_page,
|
||
"set_header_footer": _op_set_header_footer,
|
||
"remove_tables": _op_remove_tables,
|
||
}
|
||
handler = handlers.get(kind)
|
||
if handler is None:
|
||
raise ValueError(
|
||
f"不支持的操作 {kind!r}；可选：{'、'.join(sorted(handlers))}"
|
||
)
|
||
return handler(document, op)
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = SkillArgumentParser(
|
||
description="依据受控 JSON 操作编辑现有 Word DOCX 文档。"
|
||
)
|
||
parser.add_argument("--input", required=True)
|
||
parser.add_argument("--output", required=True)
|
||
parser.add_argument("--spec", help="内联 JSON 编辑说明")
|
||
parser.add_argument("--spec-file", help="JSON 编辑说明文件")
|
||
parser.add_argument("--overwrite", action="store_true")
|
||
parser.add_argument("--track-changes", action="store_true", help="将 replace_text 操作记录为修订")
|
||
parser.add_argument("--author", default="assistant", help="新修订的作者")
|
||
parser.add_argument(
|
||
"--allow-existing-revisions",
|
||
action="store_true",
|
||
help="明确允许编辑含现有修订的文档",
|
||
)
|
||
return parser
|
||
|
||
|
||
def main() -> dict[str, Any]:
|
||
from docx import Document
|
||
|
||
args = build_parser().parse_args()
|
||
source = input_file(args.input, DOCX_INPUT_SUFFIXES)
|
||
destination = output_file(
|
||
args.output,
|
||
{".docx"},
|
||
overwrite=args.overwrite,
|
||
)
|
||
if source == destination:
|
||
raise ValueError("输出路径不能与输入文件相同；请保留原始文档")
|
||
spec = load_json_argument(args.spec, args.spec_file, label="编辑说明")
|
||
unknown = set(spec) - {"operations"}
|
||
if unknown:
|
||
raise ValueError(f"编辑说明包含未知顶层字段：{sorted(unknown)}")
|
||
operations = expect_list(spec.get("operations"), "operations")
|
||
if not operations:
|
||
raise ValueError("operations 不能为空")
|
||
if len(operations) > MAX_OPERATIONS:
|
||
raise ValueError(f"operations 不能超过 {MAX_OPERATIONS} 项")
|
||
|
||
has_revisions = _has_revisions(source)
|
||
if args.track_changes and has_revisions:
|
||
raise ValueError("修订替换需要无现有修订的副本；请先根据用户要求处理原有修订")
|
||
if args.track_changes and any(expect_object(op, "operations[]").get("type") != "replace_text" for op in operations):
|
||
raise ValueError("--track-changes 仅支持 replace_text；其他操作不能伪装为修订")
|
||
if has_revisions and not args.allow_existing_revisions:
|
||
raise ValueError(
|
||
"输入文档含修订记录。为避免把未接受修订静默改坏，"
|
||
"请先使用 accept_changes.py 生成干净副本，"
|
||
"或在用户明确要求保留现有修订时传 --allow-existing-revisions"
|
||
)
|
||
|
||
document = Document(str(source))
|
||
tracker = None
|
||
if args.track_changes:
|
||
from _tracked_replace import TrackedReplacement
|
||
|
||
tracker = TrackedReplacement(source, args.author)
|
||
results: list[dict[str, Any]] = []
|
||
for index, operation in enumerate(operations):
|
||
try:
|
||
results.append(_op_replace_text(document, operation, tracker) if tracker is not None
|
||
else _apply_operation(document, operation))
|
||
except Exception as exc:
|
||
raise ValueError(f"第 {index + 1} 个操作失败：{exc}") from exc
|
||
set_update_fields(document)
|
||
|
||
descriptor, temp_name = tempfile.mkstemp(
|
||
prefix=f".{destination.stem}.",
|
||
suffix=".docx",
|
||
dir=str(destination.parent),
|
||
)
|
||
os.close(descriptor)
|
||
temp_path = Path(temp_name)
|
||
try:
|
||
document.save(str(temp_path))
|
||
archive = inspect_archive(temp_path)
|
||
if archive["missing_required_parts"]:
|
||
raise ValueError(
|
||
"编辑结果缺少必要部件："
|
||
+ "、".join(archive["missing_required_parts"])
|
||
)
|
||
Document(str(temp_path))
|
||
publish_file(temp_path, destination, overwrite=args.overwrite)
|
||
finally:
|
||
if temp_path.exists():
|
||
temp_path.unlink()
|
||
|
||
return {
|
||
"path": str(destination),
|
||
"source": str(source),
|
||
"operation_count": len(operations),
|
||
"operation_results": results,
|
||
"paragraph_count": len(document.paragraphs),
|
||
"table_count": len(document.tables),
|
||
"section_count": len(document.sections),
|
||
"source_had_revisions": has_revisions,
|
||
"tracked_replacement_count": tracker.count if tracker is not None else 0,
|
||
"archive": archive,
|
||
"requires_visual_review": True,
|
||
}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(run_cli(main))
|