Improve DOCX translation handling
This commit is contained in:
437
tests/test_docx_translator.py
Normal file
437
tests/test_docx_translator.py
Normal file
@@ -0,0 +1,437 @@
|
||||
from io import BytesIO
|
||||
import unittest
|
||||
from zipfile import ZipFile
|
||||
|
||||
import docx
|
||||
from docx.oxml import OxmlElement, parse_xml
|
||||
from docx.oxml.ns import nsdecls, qn
|
||||
from docx.shared import Inches
|
||||
from lxml import etree
|
||||
|
||||
from docutranslate.ir.document import Document
|
||||
from docutranslate.translator.ai_translator.docx_translator import (
|
||||
DocxTranslator,
|
||||
DocxTranslatorConfig,
|
||||
)
|
||||
|
||||
|
||||
NS = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
|
||||
|
||||
|
||||
def _docx_bytes(document):
|
||||
stream = BytesIO()
|
||||
document.save(stream)
|
||||
return stream.getvalue()
|
||||
|
||||
|
||||
def _text_values(docx_bytes):
|
||||
with ZipFile(BytesIO(docx_bytes)) as zf:
|
||||
tree = etree.fromstring(zf.read("word/document.xml"))
|
||||
return tree.xpath("//w:t/text()", namespaces=NS)
|
||||
|
||||
|
||||
def _comments_text_values(docx_bytes):
|
||||
with ZipFile(BytesIO(docx_bytes)) as zf:
|
||||
if "word/comments.xml" not in zf.namelist():
|
||||
return []
|
||||
tree = etree.fromstring(zf.read("word/comments.xml"))
|
||||
return tree.xpath("//w:t/text()", namespaces=NS)
|
||||
|
||||
|
||||
def _header_text_values(docx_bytes):
|
||||
values = []
|
||||
with ZipFile(BytesIO(docx_bytes)) as zf:
|
||||
for name in zf.namelist():
|
||||
if name.startswith("word/header") and name.endswith(".xml"):
|
||||
tree = etree.fromstring(zf.read(name))
|
||||
values.extend(tree.xpath("//w:t/text()", namespaces=NS))
|
||||
return values
|
||||
|
||||
|
||||
def _paragraph_and_table_counts(docx_bytes):
|
||||
with ZipFile(BytesIO(docx_bytes)) as zf:
|
||||
tree = etree.fromstring(zf.read("word/document.xml"))
|
||||
return (
|
||||
len(tree.xpath("//w:p", namespaces=NS)),
|
||||
len(tree.xpath("//w:tbl", namespaces=NS)),
|
||||
)
|
||||
|
||||
|
||||
def _section_type_counts(docx_bytes):
|
||||
with ZipFile(BytesIO(docx_bytes)) as zf:
|
||||
tree = etree.fromstring(zf.read("word/document.xml"))
|
||||
counts = {}
|
||||
for sect_pr in tree.xpath("//w:sectPr", namespaces=NS):
|
||||
value = sect_pr.xpath("./w:type/@w:val", namespaces=NS)
|
||||
key = value[0] if value else "nextPage(default)"
|
||||
counts[key] = counts.get(key, 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
def _add_toc_paragraph(document):
|
||||
paragraph = document.add_paragraph()
|
||||
p_style = OxmlElement("w:pStyle")
|
||||
p_style.set(qn("w:val"), "TOC1")
|
||||
paragraph._p.get_or_add_pPr().append(p_style)
|
||||
|
||||
hyperlink = OxmlElement("w:hyperlink")
|
||||
hyperlink.set(qn("w:anchor"), "_TocFixture")
|
||||
|
||||
title_run = OxmlElement("w:r")
|
||||
title_text = OxmlElement("w:t")
|
||||
title_text.text = "Purpose"
|
||||
title_run.append(title_text)
|
||||
|
||||
tab_run = OxmlElement("w:r")
|
||||
tab_run.append(OxmlElement("w:tab"))
|
||||
|
||||
page_run = OxmlElement("w:r")
|
||||
page_text = OxmlElement("w:t")
|
||||
page_text.text = "3"
|
||||
page_run.append(page_text)
|
||||
|
||||
hyperlink.extend([title_run, tab_run, page_run])
|
||||
paragraph._p.append(hyperlink)
|
||||
|
||||
|
||||
def _add_internal_hyperlink_paragraph(document):
|
||||
paragraph = document.add_paragraph()
|
||||
paragraph.add_run("Read ")
|
||||
hyperlink = OxmlElement("w:hyperlink")
|
||||
hyperlink.set(qn("w:anchor"), "_Manual")
|
||||
link_run = OxmlElement("w:r")
|
||||
link_text = OxmlElement("w:t")
|
||||
link_text.text = "manual"
|
||||
link_run.append(link_text)
|
||||
hyperlink.append(link_run)
|
||||
paragraph._p.append(hyperlink)
|
||||
paragraph.add_run(" now")
|
||||
|
||||
|
||||
def _add_textbox_paragraph(document):
|
||||
paragraph = document.add_paragraph()
|
||||
paragraph.add_run("Diagram: ")
|
||||
textbox_run = parse_xml(
|
||||
f"""
|
||||
<w:r {nsdecls("w", "wp", "a")} xmlns:wps="http://schemas.microsoft.com/office/word/2010/wordprocessingShape">
|
||||
<w:drawing>
|
||||
<wp:inline>
|
||||
<a:graphic>
|
||||
<a:graphicData uri="http://schemas.microsoft.com/office/word/2010/wordprocessingShape">
|
||||
<wps:wsp>
|
||||
<wps:txbx>
|
||||
<w:txbxContent>
|
||||
<w:p>
|
||||
<w:r><w:t>Flow label</w:t></w:r>
|
||||
</w:p>
|
||||
</w:txbxContent>
|
||||
</wps:txbx>
|
||||
</wps:wsp>
|
||||
</a:graphicData>
|
||||
</a:graphic>
|
||||
</wp:inline>
|
||||
</w:drawing>
|
||||
</w:r>
|
||||
"""
|
||||
)
|
||||
paragraph._p.append(textbox_run)
|
||||
|
||||
|
||||
def _add_inline_sdt(paragraph, text):
|
||||
sdt = OxmlElement("w:sdt")
|
||||
sdt_content = OxmlElement("w:sdtContent")
|
||||
run = OxmlElement("w:r")
|
||||
text_element = OxmlElement("w:t")
|
||||
text_element.text = text
|
||||
run.append(text_element)
|
||||
sdt_content.append(run)
|
||||
sdt.append(sdt_content)
|
||||
paragraph._p.append(sdt)
|
||||
|
||||
|
||||
def _add_block_sdt(document, text):
|
||||
sdt = OxmlElement("w:sdt")
|
||||
sdt_content = OxmlElement("w:sdtContent")
|
||||
paragraph = OxmlElement("w:p")
|
||||
run = OxmlElement("w:r")
|
||||
text_element = OxmlElement("w:t")
|
||||
text_element.text = text
|
||||
run.append(text_element)
|
||||
paragraph.append(run)
|
||||
sdt_content.append(paragraph)
|
||||
sdt.append(sdt_content)
|
||||
document._body._element.append(sdt)
|
||||
|
||||
|
||||
def _add_page_field_paragraph(document):
|
||||
paragraph = document.add_paragraph("Page ")
|
||||
for fld_type in ("begin", "separate", "end"):
|
||||
run = OxmlElement("w:r")
|
||||
if fld_type == "begin":
|
||||
fld = OxmlElement("w:fldChar")
|
||||
fld.set(qn("w:fldCharType"), "begin")
|
||||
run.append(fld)
|
||||
instr_run = OxmlElement("w:r")
|
||||
instr = OxmlElement("w:instrText")
|
||||
instr.text = " PAGE "
|
||||
instr_run.append(instr)
|
||||
paragraph._p.append(run)
|
||||
paragraph._p.append(instr_run)
|
||||
continue
|
||||
if fld_type == "separate":
|
||||
fld = OxmlElement("w:fldChar")
|
||||
fld.set(qn("w:fldCharType"), "separate")
|
||||
run.append(fld)
|
||||
cached_run = OxmlElement("w:r")
|
||||
cached_text = OxmlElement("w:t")
|
||||
cached_text.text = "1"
|
||||
cached_run.append(cached_text)
|
||||
paragraph._p.append(run)
|
||||
paragraph._p.append(cached_run)
|
||||
continue
|
||||
fld = OxmlElement("w:fldChar")
|
||||
fld.set(qn("w:fldCharType"), "end")
|
||||
run.append(fld)
|
||||
paragraph._p.append(run)
|
||||
|
||||
|
||||
def _add_tracked_changes_paragraph(document):
|
||||
paragraph = document.add_paragraph("Revision ")
|
||||
|
||||
deleted = OxmlElement("w:del")
|
||||
deleted.set(qn("w:id"), "1")
|
||||
deleted_run = OxmlElement("w:r")
|
||||
deleted_text = OxmlElement("w:delText")
|
||||
deleted_text.text = "old text"
|
||||
deleted_run.append(deleted_text)
|
||||
deleted.append(deleted_run)
|
||||
paragraph._p.append(deleted)
|
||||
|
||||
inserted = OxmlElement("w:ins")
|
||||
inserted.set(qn("w:id"), "2")
|
||||
inserted_run = OxmlElement("w:r")
|
||||
inserted_text = OxmlElement("w:t")
|
||||
inserted_text.text = "new text"
|
||||
inserted_run.append(inserted_text)
|
||||
inserted.append(inserted_run)
|
||||
paragraph._p.append(inserted)
|
||||
|
||||
|
||||
def _add_redundant_section_break(document):
|
||||
document.add_paragraph("Section One")
|
||||
for _ in range(2):
|
||||
section = document.add_section()
|
||||
section.start_type = docx.enum.section.WD_SECTION.NEW_PAGE
|
||||
section.top_margin = Inches(1)
|
||||
section.bottom_margin = Inches(1)
|
||||
section.left_margin = Inches(1)
|
||||
section.right_margin = Inches(1)
|
||||
document.add_paragraph("Section Two")
|
||||
|
||||
|
||||
def _build_fixture_docx():
|
||||
document = docx.Document()
|
||||
for section in document.sections:
|
||||
section.top_margin = Inches(1)
|
||||
section.bottom_margin = Inches(1)
|
||||
section.left_margin = Inches(1)
|
||||
section.right_margin = Inches(1)
|
||||
document.add_heading("Fixture", level=1)
|
||||
document.add_paragraph("Normal paragraph")
|
||||
_add_toc_paragraph(document)
|
||||
_add_internal_hyperlink_paragraph(document)
|
||||
_add_textbox_paragraph(document)
|
||||
inline_sdt_paragraph = document.add_paragraph("Controlled ")
|
||||
_add_inline_sdt(inline_sdt_paragraph, "Approval Required")
|
||||
_add_block_sdt(document, "Block Controlled Text")
|
||||
_add_page_field_paragraph(document)
|
||||
_add_tracked_changes_paragraph(document)
|
||||
|
||||
table = document.add_table(rows=2, cols=2)
|
||||
merged = table.cell(0, 0).merge(table.cell(0, 1))
|
||||
merged.text = "Merged Cell"
|
||||
table.cell(1, 0).text = "Left Cell"
|
||||
table.cell(1, 1).text = "Right Cell"
|
||||
nested_table = table.cell(1, 0).add_table(rows=1, cols=1)
|
||||
nested_table.cell(0, 0).text = "Nested Cell"
|
||||
|
||||
comment_anchor = document.add_paragraph().add_run("Comment anchor")
|
||||
document.add_comment(comment_anchor, text="Review this note", author="QA")
|
||||
|
||||
_add_redundant_section_break(document)
|
||||
for sect_pr in document._element.xpath(".//w:sectPr"):
|
||||
for child in list(sect_pr):
|
||||
if child.tag in {qn("w:headerReference"), qn("w:footerReference")}:
|
||||
sect_pr.remove(child)
|
||||
return _docx_bytes(document)
|
||||
|
||||
|
||||
class DocxTranslatorTest(unittest.TestCase):
|
||||
def _translator(self, skip_translate=True):
|
||||
return DocxTranslator(DocxTranslatorConfig(skip_translate=skip_translate))
|
||||
|
||||
def test_pre_translate_extracts_docx_boundary_text(self):
|
||||
translator = self._translator()
|
||||
_, _, originals = translator._pre_translate(Document(content=_build_fixture_docx(), suffix=".docx"))
|
||||
|
||||
self.assertIn("Purpose", originals)
|
||||
self.assertIn("Flow label", originals)
|
||||
self.assertIn("Controlled Approval Required", originals)
|
||||
self.assertIn("Block Controlled Text", originals)
|
||||
self.assertIn("Revision new text", originals)
|
||||
self.assertNotIn("Revision old textnew text", originals)
|
||||
self.assertTrue(all("old text" not in original for original in originals))
|
||||
self.assertIn("Merged Cell", originals)
|
||||
self.assertIn("Nested Cell", originals)
|
||||
self.assertIn("Review this note", originals)
|
||||
self.assertNotIn("Purpose3", originals)
|
||||
self.assertNotIn("PAGE", originals)
|
||||
self.assertIn("Page ", originals)
|
||||
self.assertNotIn("Page 1", originals)
|
||||
|
||||
def test_replace_mode_writes_boundary_text_without_structural_loss(self):
|
||||
input_bytes = _build_fixture_docx()
|
||||
translator = self._translator()
|
||||
doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
|
||||
translations = {
|
||||
"Purpose": "Tujuan",
|
||||
"Flow label": "Label alur",
|
||||
"Controlled Approval Required": "Persetujuan diperlukan",
|
||||
"Block Controlled Text": "Teks kontrol blok",
|
||||
"Revision new text": "Revisi teks baru",
|
||||
"Merged Cell": "Sel gabungan",
|
||||
"Nested Cell": "Sel bersarang",
|
||||
"Review this note": "Tinjau catatan ini",
|
||||
}
|
||||
translated = [translations.get(text, text) for text in originals]
|
||||
|
||||
output_bytes = translator._after_translate(doc, elements, translated, originals)
|
||||
|
||||
body_text = _text_values(output_bytes)
|
||||
comment_text = _comments_text_values(output_bytes)
|
||||
self.assertIn("Tujuan", body_text)
|
||||
self.assertIn("3", body_text)
|
||||
self.assertIn("Label alur", body_text)
|
||||
self.assertIn("Persetujuan diperlukan", "".join(body_text))
|
||||
self.assertIn("Teks kontrol blok", body_text)
|
||||
self.assertIn("Revisi teks baru", "".join(body_text))
|
||||
self.assertIn("Sel gabungan", body_text)
|
||||
self.assertIn("Sel bersarang", body_text)
|
||||
self.assertIn("Tinjau catatan ini", comment_text)
|
||||
self.assertEqual(_paragraph_and_table_counts(input_bytes), _paragraph_and_table_counts(output_bytes))
|
||||
|
||||
def test_append_mode_keeps_textbox_mapping(self):
|
||||
input_bytes = _build_fixture_docx()
|
||||
translator = DocxTranslator(DocxTranslatorConfig(skip_translate=True, insert_mode="append"))
|
||||
doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
|
||||
translated = ["Label alur" if text == "Flow label" else text for text in originals]
|
||||
|
||||
output_bytes = translator._after_translate(doc, elements, translated, originals)
|
||||
|
||||
body_text = _text_values(output_bytes)
|
||||
self.assertIn("Flow label", body_text)
|
||||
self.assertIn("Label alur", body_text)
|
||||
|
||||
def test_real_translation_mode_normalizes_redundant_section_breaks(self):
|
||||
input_bytes = _build_fixture_docx()
|
||||
translator = DocxTranslator(
|
||||
DocxTranslatorConfig(
|
||||
skip_translate=False,
|
||||
api_key="dummy",
|
||||
base_url="http://127.0.0.1",
|
||||
model_id="dummy",
|
||||
)
|
||||
)
|
||||
doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
|
||||
for sect_pr in doc._element.xpath(".//w:sectPr"):
|
||||
for child in list(sect_pr):
|
||||
if child.tag in {qn("w:headerReference"), qn("w:footerReference")}:
|
||||
sect_pr.remove(child)
|
||||
|
||||
output_bytes = translator._after_translate(doc, elements, originals, originals)
|
||||
|
||||
source_counts = _section_type_counts(input_bytes)
|
||||
output_counts = _section_type_counts(output_bytes)
|
||||
self.assertGreater(source_counts.get("nextPage(default)", 0), output_counts.get("nextPage(default)", 0))
|
||||
self.assertGreater(output_counts.get("continuous", 0), source_counts.get("continuous", 0))
|
||||
|
||||
def test_hyperlink_runs_are_not_merged_across_parent_boundary(self):
|
||||
input_bytes = _build_fixture_docx()
|
||||
translator = self._translator()
|
||||
doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
|
||||
translated = ["Baca manual sekarang" if text == "Read manual now" else text for text in originals]
|
||||
|
||||
output_bytes = translator._after_translate(doc, elements, translated, originals)
|
||||
|
||||
with ZipFile(BytesIO(output_bytes)) as zf:
|
||||
tree = etree.fromstring(zf.read("word/document.xml"))
|
||||
self.assertEqual(len(tree.xpath("//w:hyperlink", namespaces=NS)), 2)
|
||||
|
||||
def test_shared_headers_are_extracted_once_across_sections(self):
|
||||
document = docx.Document()
|
||||
document.sections[0].header.paragraphs[0].text = "Shared Header"
|
||||
document.add_paragraph("Body One")
|
||||
document.add_section()
|
||||
document.add_paragraph("Body Two")
|
||||
input_bytes = _docx_bytes(document)
|
||||
|
||||
translator = self._translator()
|
||||
_, _, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
|
||||
|
||||
self.assertEqual(originals.count("Shared Header"), 1)
|
||||
|
||||
def test_shared_header_is_written_once_and_preserved(self):
|
||||
document = docx.Document()
|
||||
document.sections[0].header.paragraphs[0].text = "Shared Header"
|
||||
document.add_paragraph("Body One")
|
||||
document.add_section()
|
||||
document.add_paragraph("Body Two")
|
||||
input_bytes = _docx_bytes(document)
|
||||
|
||||
translator = self._translator()
|
||||
doc, elements, originals = translator._pre_translate(Document(content=input_bytes, suffix=".docx"))
|
||||
translated = ["Header Bersama" if text == "Shared Header" else text for text in originals]
|
||||
output_bytes = translator._after_translate(doc, elements, translated, originals)
|
||||
|
||||
self.assertEqual(_header_text_values(output_bytes).count("Header Bersama"), 1)
|
||||
|
||||
def test_translate_method_end_to_end_with_fake_agent(self):
|
||||
class FakeAgent:
|
||||
def send_segments(self, segments, chunk_size):
|
||||
translations = {
|
||||
"Purpose": "Tujuan",
|
||||
"Flow label": "Label alur",
|
||||
"Controlled Approval Required": "Persetujuan diperlukan",
|
||||
"Block Controlled Text": "Teks kontrol blok",
|
||||
"Revision new text": "Revisi teks baru",
|
||||
"Nested Cell": "Sel bersarang",
|
||||
"Review this note": "Tinjau catatan ini",
|
||||
}
|
||||
return [translations.get(segment, segment) for segment in segments]
|
||||
|
||||
document = Document(content=_build_fixture_docx(), suffix=".docx")
|
||||
translator = DocxTranslator(
|
||||
DocxTranslatorConfig(
|
||||
skip_translate=False,
|
||||
api_key="dummy",
|
||||
base_url="http://127.0.0.1",
|
||||
model_id="dummy",
|
||||
)
|
||||
)
|
||||
translator.translate_agent = FakeAgent()
|
||||
|
||||
translator.translate(document)
|
||||
|
||||
body_text = _text_values(document.content)
|
||||
comments_text = _comments_text_values(document.content)
|
||||
self.assertIn("Tujuan", body_text)
|
||||
self.assertIn("Label alur", body_text)
|
||||
self.assertIn("Persetujuan diperlukan", "".join(body_text))
|
||||
self.assertIn("Teks kontrol blok", body_text)
|
||||
self.assertIn("Revisi teks baru", "".join(body_text))
|
||||
self.assertIn("Sel bersarang", body_text)
|
||||
self.assertIn("Tinjau catatan ini", comments_text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user