391 lines
13 KiB
Python
391 lines
13 KiB
Python
from io import BytesIO
|
|
from copy import deepcopy
|
|
|
|
from docx import Document
|
|
from docx.document import Document as DocumentObject
|
|
from docx.oxml import OxmlElement
|
|
from docx.oxml.table import CT_Tbl
|
|
from docx.oxml.text.paragraph import CT_P
|
|
from docx.table import Table, _Cell
|
|
from docx.text.paragraph import Paragraph
|
|
|
|
|
|
def _iter_block_items(parent: DocumentObject | _Cell):
|
|
parent_elm = parent.element.body if isinstance(parent, DocumentObject) else parent._tc
|
|
for child in parent_elm.iterchildren():
|
|
if isinstance(child, CT_P):
|
|
yield Paragraph(child, parent)
|
|
elif isinstance(child, CT_Tbl):
|
|
yield Table(child, parent)
|
|
|
|
|
|
def _is_heading(paragraph: Paragraph) -> bool:
|
|
style_name = paragraph.style.name if paragraph.style is not None else ""
|
|
normalized = style_name.lower().replace(" ", "")
|
|
return normalized.startswith("heading")
|
|
|
|
|
|
def _delete_block(block):
|
|
element = block._element
|
|
parent = element.getparent()
|
|
if parent is not None:
|
|
parent.remove(element)
|
|
|
|
|
|
def _delete_heading_section(heading: Paragraph):
|
|
blocks = [heading]
|
|
current = heading._element.getnext()
|
|
while current is not None:
|
|
if isinstance(current, CT_P):
|
|
para = Paragraph(current, heading._parent)
|
|
if _is_heading(para):
|
|
break
|
|
blocks.append(para)
|
|
elif isinstance(current, CT_Tbl):
|
|
blocks.append(Table(current, heading._parent))
|
|
current = current.getnext()
|
|
for block in blocks:
|
|
_delete_block(block)
|
|
|
|
|
|
def _remove_unreferenced_headings(document: DocumentObject, referenced_anchors: set[str]):
|
|
headings_to_remove: list[Paragraph] = []
|
|
found_first_heading = False
|
|
pre_heading_blocks: list = []
|
|
print(f"[EXPORT] referenced_anchors: {referenced_anchors}")
|
|
for block in _iter_block_items(document):
|
|
if isinstance(block, Paragraph) and _is_heading(block):
|
|
found_first_heading = True
|
|
text = block.text.strip()
|
|
if text not in referenced_anchors:
|
|
print(f"[EXPORT] Unreferenced heading found, will remove: '{text}'")
|
|
headings_to_remove.append(block)
|
|
elif not found_first_heading:
|
|
pre_heading_blocks.append(block)
|
|
for heading in headings_to_remove:
|
|
_delete_heading_section(heading)
|
|
if not referenced_anchors:
|
|
for block in pre_heading_blocks:
|
|
_delete_block(block)
|
|
|
|
|
|
def _clear_paragraph(paragraph: Paragraph):
|
|
element = paragraph._element
|
|
for child in list(element):
|
|
if child.tag.endswith("}r"):
|
|
element.remove(child)
|
|
|
|
|
|
def _copy_paragraph_format(target: Paragraph, source: Paragraph | None):
|
|
if source is None:
|
|
return
|
|
source_ppr = source._element.pPr
|
|
if source_ppr is not None:
|
|
target._element.insert(0, deepcopy(source_ppr))
|
|
|
|
|
|
def _copy_run_format(target_run, source_paragraph: Paragraph | None):
|
|
if source_paragraph is None:
|
|
return
|
|
for source_run in source_paragraph.runs:
|
|
if source_run._element.rPr is not None:
|
|
target_run._element.insert(0, deepcopy(source_run._element.rPr))
|
|
break
|
|
|
|
|
|
def _extract_first_run_format(source_paragraph: Paragraph | None):
|
|
if source_paragraph is None:
|
|
return None
|
|
for source_run in source_paragraph.runs:
|
|
if source_run._element.rPr is not None:
|
|
return deepcopy(source_run._element.rPr)
|
|
return None
|
|
|
|
|
|
def _set_paragraph_text(
|
|
paragraph: Paragraph,
|
|
text: str,
|
|
style_name: str | None = None,
|
|
template_paragraph: Paragraph | None = None,
|
|
):
|
|
run_format = _extract_first_run_format(template_paragraph)
|
|
_clear_paragraph(paragraph)
|
|
if style_name:
|
|
try:
|
|
paragraph.style = style_name
|
|
except Exception:
|
|
pass
|
|
if text:
|
|
run = paragraph.add_run(text)
|
|
if run_format is not None:
|
|
run._element.insert(0, run_format)
|
|
|
|
|
|
def _append_paragraph_after(
|
|
paragraph: Paragraph,
|
|
text: str,
|
|
style_name: str | None = None,
|
|
template_paragraph: Paragraph | None = None,
|
|
) -> Paragraph:
|
|
new_p = OxmlElement("w:p")
|
|
paragraph._element.addnext(new_p)
|
|
new_para = Paragraph(new_p, paragraph._parent)
|
|
_copy_paragraph_format(new_para, template_paragraph)
|
|
if style_name:
|
|
try:
|
|
new_para.style = style_name
|
|
except Exception:
|
|
pass
|
|
if text:
|
|
run = new_para.add_run(text)
|
|
_copy_run_format(run, template_paragraph)
|
|
return new_para
|
|
|
|
|
|
def _set_cell_text_with_template(cell, value: str, template_paragraph: Paragraph | None = None):
|
|
if not cell.paragraphs:
|
|
cell.text = value
|
|
return
|
|
paragraph = cell.paragraphs[0]
|
|
_clear_paragraph(paragraph)
|
|
run = paragraph.add_run(value)
|
|
_copy_run_format(run, template_paragraph)
|
|
|
|
|
|
def _resize_table_rows(table: Table, row_count: int):
|
|
current_rows = len(table.rows)
|
|
if current_rows == 0:
|
|
return
|
|
if current_rows < row_count:
|
|
template_row = table.rows[-1]._tr
|
|
for _ in range(row_count - current_rows):
|
|
table._tbl.append(deepcopy(template_row))
|
|
elif current_rows > row_count:
|
|
for _ in range(current_rows - row_count):
|
|
table._tbl.remove(table.rows[-1]._tr)
|
|
|
|
|
|
def _fill_table(table: Table, matrix: list[list[str]]):
|
|
if not matrix:
|
|
return
|
|
_resize_table_rows(table, len(matrix))
|
|
template_cell_paragraph = table.rows[0].cells[0].paragraphs[0] if table.rows and table.rows[0].cells else None
|
|
for row_index, row_values in enumerate(matrix):
|
|
row = table.rows[row_index]
|
|
for col_index, cell in enumerate(row.cells):
|
|
value = row_values[col_index] if col_index < len(row_values) else ""
|
|
_set_cell_text_with_template(cell, value, template_cell_paragraph)
|
|
|
|
|
|
def _append_table_after(
|
|
paragraph: Paragraph,
|
|
rows: list[list[str]],
|
|
headers: list[str] | None = None,
|
|
template_table: Table | None = None,
|
|
):
|
|
matrix = [headers, *rows] if headers else rows
|
|
if template_table is not None:
|
|
cloned_tbl = deepcopy(template_table._tbl)
|
|
paragraph._element.addnext(cloned_tbl)
|
|
cloned_table = Table(cloned_tbl, paragraph._parent)
|
|
_fill_table(cloned_table, matrix)
|
|
return cloned_table
|
|
|
|
container = paragraph._parent
|
|
table = container.add_table(rows=max(len(matrix), 1), cols=max(len(headers or []), len(rows[0]) if rows else 1))
|
|
if headers:
|
|
for row_index, row_values in enumerate(matrix):
|
|
for index, value in enumerate(row_values):
|
|
table.rows[row_index].cells[index].text = value
|
|
elif rows:
|
|
for row_index, row_values in enumerate(matrix):
|
|
for index, value in enumerate(row_values):
|
|
table.rows[row_index].cells[index].text = value
|
|
|
|
tbl = table._tbl
|
|
tbl.getparent().remove(tbl)
|
|
paragraph._element.addnext(tbl)
|
|
return Table(tbl, container)
|
|
|
|
|
|
def _append_empty_paragraph_after_table(table: Table, style_name: str | None = None) -> Paragraph:
|
|
new_p = OxmlElement("w:p")
|
|
table._tbl.addnext(new_p)
|
|
new_para = Paragraph(new_p, table._parent)
|
|
if style_name:
|
|
try:
|
|
new_para.style = style_name
|
|
except Exception:
|
|
pass
|
|
return new_para
|
|
|
|
|
|
def _find_heading_paragraph(document: DocumentObject, heading_text: str, after_element=None) -> Paragraph | None:
|
|
started = after_element is None
|
|
for block in _iter_block_items(document):
|
|
if isinstance(block, Paragraph) and _is_heading(block) and block.text.strip() == heading_text.strip():
|
|
if started:
|
|
return block
|
|
if after_element is not None and block._element == after_element:
|
|
started = True
|
|
return None
|
|
|
|
|
|
def _collect_section_templates(heading: Paragraph):
|
|
first_body_style = None
|
|
paragraph_template = None
|
|
table_template = None
|
|
blocks = []
|
|
current = heading._element.getnext()
|
|
while current is not None:
|
|
if isinstance(current, CT_P):
|
|
current_paragraph = Paragraph(current, heading._parent)
|
|
if _is_heading(current_paragraph):
|
|
break
|
|
if first_body_style is None and current_paragraph.style is not None:
|
|
first_body_style = current_paragraph.style.name
|
|
if paragraph_template is None:
|
|
paragraph_template = current_paragraph
|
|
blocks.append(current_paragraph)
|
|
elif isinstance(current, CT_Tbl):
|
|
current_table = Table(current, heading._parent)
|
|
if table_template is None:
|
|
table_template = current_table
|
|
blocks.append(current_table)
|
|
current = current.getnext()
|
|
return first_body_style, paragraph_template, table_template, blocks
|
|
|
|
|
|
def _insert_content_after(
|
|
insert_after: Paragraph,
|
|
content: dict,
|
|
first_body_style: str | None,
|
|
paragraph_template: Paragraph | None,
|
|
table_template: Table | None,
|
|
):
|
|
current_anchor: Paragraph = insert_after
|
|
content_blocks = content.get("content", [])
|
|
for block in content_blocks:
|
|
block_type = block.get("type")
|
|
if block_type == "table":
|
|
rows = [list(row) for row in block.get("rows", [])]
|
|
headers = block.get("headers") or []
|
|
table = _append_table_after(current_anchor, rows, headers, table_template)
|
|
current_anchor = _append_empty_paragraph_after_table(table, first_body_style)
|
|
else:
|
|
text = block.get("text", "")
|
|
text_parts = [item for item in text.split("\n") if item] or [text]
|
|
for text_part in text_parts:
|
|
current_anchor = _append_paragraph_after(
|
|
current_anchor,
|
|
text_part,
|
|
first_body_style,
|
|
paragraph_template,
|
|
)
|
|
return current_anchor
|
|
|
|
|
|
def _replace_section_content(
|
|
document: DocumentObject,
|
|
anchor_title: str,
|
|
target_title: str,
|
|
content: dict,
|
|
write_mode: str,
|
|
after_element=None,
|
|
):
|
|
heading = _find_heading_paragraph(document, anchor_title, after_element)
|
|
if heading is None:
|
|
return after_element
|
|
_set_paragraph_text(heading, target_title, heading.style.name if heading.style is not None else None, heading)
|
|
first_body_style, paragraph_template, table_template, blocks_to_remove = _collect_section_templates(heading)
|
|
|
|
if write_mode == "replace_heading_only":
|
|
return heading._element
|
|
|
|
if write_mode == "replace_section":
|
|
for block in blocks_to_remove:
|
|
_delete_block(block)
|
|
|
|
_insert_content_after(
|
|
heading,
|
|
content,
|
|
first_body_style,
|
|
paragraph_template,
|
|
table_template,
|
|
)
|
|
return heading._element
|
|
|
|
|
|
def _group_logs(logs: list[dict]) -> list[list[dict]]:
|
|
groups: list[list[dict]] = []
|
|
for item in logs:
|
|
anchor_title = item.get("anchor_title") or item.get("title") or ""
|
|
if not groups:
|
|
groups.append([item])
|
|
continue
|
|
last_group = groups[-1]
|
|
last_anchor = last_group[0].get("anchor_title") or last_group[0].get("title") or ""
|
|
if anchor_title == last_anchor:
|
|
last_group.append(item)
|
|
else:
|
|
groups.append([item])
|
|
return groups
|
|
|
|
|
|
def _replace_section_group(
|
|
document: DocumentObject,
|
|
items: list[dict],
|
|
after_element=None,
|
|
):
|
|
first_item = items[0]
|
|
anchor_title = first_item.get("anchor_title") or first_item.get("title") or ""
|
|
target_title = first_item.get("title") or anchor_title
|
|
heading = _find_heading_paragraph(document, anchor_title, after_element)
|
|
if heading is None:
|
|
return after_element
|
|
|
|
_set_paragraph_text(heading, target_title, heading.style.name if heading.style is not None else None, heading)
|
|
first_body_style, paragraph_template, table_template, blocks_to_remove = _collect_section_templates(heading)
|
|
|
|
if len(items) == 1 and first_item.get("write_mode") == "replace_heading_only":
|
|
return heading._element
|
|
|
|
preserve_existing = len(items) == 1 and first_item.get("write_mode") == "append_after_heading"
|
|
if not preserve_existing:
|
|
for block in blocks_to_remove:
|
|
_delete_block(block)
|
|
|
|
current_anchor = heading
|
|
for item in items:
|
|
current_anchor = _insert_content_after(
|
|
current_anchor,
|
|
item.get("content") or {"content": []},
|
|
first_body_style,
|
|
paragraph_template,
|
|
table_template,
|
|
)
|
|
return heading._element
|
|
|
|
|
|
def export_document_bytes(template_bytes: bytes, logs: list[dict]) -> bytes:
|
|
document = Document(BytesIO(template_bytes))
|
|
|
|
referenced_anchors: set[str] = set()
|
|
for item in logs:
|
|
for key in ("anchor_title", "title"):
|
|
val = (item.get(key) or "").strip()
|
|
if val:
|
|
referenced_anchors.add(val)
|
|
|
|
print(f"[EXPORT] logs count={len(logs)}, anchor_titles={[(l.get('anchor_title'), l.get('title')) for l in logs]}")
|
|
|
|
last_heading_element = None
|
|
for group in _group_logs(logs):
|
|
last_heading_element = _replace_section_group(document, group, last_heading_element)
|
|
|
|
_remove_unreferenced_headings(document, referenced_anchors)
|
|
|
|
output = BytesIO()
|
|
document.save(output)
|
|
return output.getvalue()
|