zwt13703
142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
from io import BytesIO
|
|
|
|
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 _clear_paragraph(paragraph: Paragraph):
|
|
element = paragraph._element
|
|
for child in list(element):
|
|
if child.tag.endswith("}r"):
|
|
element.remove(child)
|
|
|
|
|
|
def _append_paragraph_after(paragraph: Paragraph, text: str, style_name: str | None = None) -> Paragraph:
|
|
new_p = OxmlElement("w:p")
|
|
paragraph._element.addnext(new_p)
|
|
new_para = Paragraph(new_p, paragraph._parent)
|
|
if style_name:
|
|
try:
|
|
new_para.style = style_name
|
|
except Exception:
|
|
pass
|
|
if text:
|
|
new_para.add_run(text)
|
|
return new_para
|
|
|
|
|
|
def _append_table_after(paragraph: Paragraph, rows: list[list[str]], headers: list[str] | None = None):
|
|
container = paragraph._parent
|
|
table = container.add_table(rows=1, cols=max(len(headers or []), len(rows[0]) if rows else 1))
|
|
if headers:
|
|
header_cells = table.rows[0].cells
|
|
for index, value in enumerate(headers):
|
|
header_cells[index].text = value
|
|
else:
|
|
if rows:
|
|
first = rows.pop(0)
|
|
for index, value in enumerate(first):
|
|
table.rows[0].cells[index].text = value
|
|
for row in rows:
|
|
new_row = table.add_row().cells
|
|
for index, value in enumerate(row):
|
|
new_row[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) -> Paragraph | None:
|
|
for block in _iter_block_items(document):
|
|
if isinstance(block, Paragraph) and _is_heading(block) and block.text.strip() == heading_text.strip():
|
|
return block
|
|
return None
|
|
|
|
|
|
def _replace_section_content(document: DocumentObject, heading_title: str, content: dict):
|
|
heading = _find_heading_paragraph(document, heading_title)
|
|
if heading is None:
|
|
return
|
|
|
|
first_body_style = None
|
|
current = heading._element.getnext()
|
|
blocks_to_remove = []
|
|
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
|
|
blocks_to_remove.append(current_paragraph)
|
|
elif isinstance(current, CT_Tbl):
|
|
blocks_to_remove.append(Table(current, heading._parent))
|
|
current = current.getnext()
|
|
|
|
for block in blocks_to_remove:
|
|
_delete_block(block)
|
|
|
|
insert_after = heading
|
|
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(insert_after, rows, headers)
|
|
insert_after = _append_empty_paragraph_after_table(table, first_body_style)
|
|
else:
|
|
text = block.get("text", "")
|
|
insert_after = _append_paragraph_after(insert_after, text, first_body_style)
|
|
|
|
|
|
def export_document_bytes(template_bytes: bytes, logs: list[dict]) -> bytes:
|
|
document = Document(BytesIO(template_bytes))
|
|
for item in logs:
|
|
_replace_section_content(document, item["title"], item["content"])
|
|
|
|
output = BytesIO()
|
|
document.save(output)
|
|
return output.getvalue()
|