proposa-ai/pdf_table_to_docx/docx_writer.py

"""
docx_writer.py - DOCX 文档生成模块

负责：
- 使用 python-docx 创建 Word 文档
- 根据 TableInfo 在文档中重建表格
- 设置合并单元格（跨行/跨列）
- 设置单元格对齐、边框等基本格式
"""

import logging
from typing import Optional

from docx import Document
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.ns import qn, nsdecls
from docx.oxml import parse_xml
from docx.shared import Emu, Pt, Cm
from docx.table import _Cell

from .table_parser import TableInfo

logger = logging.getLogger(__name__)


def _get_alignment(text: str) -> int:
    """从文本内容中猜测对齐方式。"""
    if not text or not text.strip():
        return WD_ALIGN_PARAGRAPH.LEFT
    stripped = text.strip()
    if stripped.replace(",", "").replace(".", "").replace("-", "").isdigit():
        return WD_ALIGN_PARAGRAPH.RIGHT
    if len(stripped) <= 5:
        return WD_ALIGN_PARAGRAPH.CENTER
    return WD_ALIGN_PARAGRAPH.LEFT


def _convert_cm_to_emu(cm_value: float) -> int:
    """将厘米转换为 EMU（docx 内部单位）。"""
    return int(round(cm_value * 360000))


def _merge_cells(table, r: int, c: int, row_span: int, col_span: int):
    """合并表格中的单元格。"""
    if row_span <= 1 and col_span <= 1:
        return
    end_r = r + row_span - 1
    end_c = c + col_span - 1
    start_cell = table.cell(r, c)
    end_cell = table.cell(end_r, end_c)
    start_cell.merge(end_cell)


def write_table_to_doc(doc, table_info: TableInfo, body_width_cm: float):
    """将单个 TableInfo 写入到已打开的 doc 对象中。

    Args:
        doc: Document 对象
        table_info: 要写入的表格
        body_width_cm: 正文宽度（厘米），用于计算列宽
    """
    if table_info.rows == 0 or table_info.cols == 0:
        return

    doc_table = doc.add_table(
        rows=table_info.rows,
        cols=table_info.cols,
    )
    doc_table.autofit = False
    doc_table.allow_autofit = False

    # ---- 设置列宽 ----
    usable_width = body_width_cm * 0.95
    if table_info.col_widths:
        for c_idx, ratio in enumerate(table_info.col_widths):
            width_emu = _convert_cm_to_emu(usable_width * ratio)
            for row in doc_table.rows:
                row.cells[c_idx].width = Emu(width_emu)

    # ---- 填充内容、合并单元格 ----
    for r in range(table_info.rows):
        for c in range(table_info.cols):
            cell_info = table_info.cells[r][c]

            if cell_info.row_span == 0 or cell_info.col_span == 0:
                continue

            _merge_cells(doc_table, r, c, cell_info.row_span, cell_info.col_span)
            cell = doc_table.cell(r, c)

            text = cell_info.text or ""
            cell.text = ""

            lines = text.split("\n")
            for li, line in enumerate(lines):
                if li > 0:
                    cell.add_paragraph()
                run = cell.paragraphs[li].add_run(line.strip())
                run.font.size = Pt(10)
                run.font.name = "宋体"
                run._element.rPr.rFonts.set(qn("w:eastAsia"), "宋体")

            alignment = _get_alignment(text)
            for para in cell.paragraphs:
                para.alignment = alignment

            cell.vertical_alignment = 1  # CENTER

    # ---- 设置表格边框 ----
    tbl = doc_table._tbl
    tblPr = (
        tbl.tblPr
        if tbl.tblPr is not None
        else parse_xml(f'<w:tblPr {nsdecls("w")}/>')
    )

    borders = parse_xml(
        f'<w:tblBorders {nsdecls("w")}>'
        '  <w:top w:val="single" w:sz="4" w:space="0" w:color="000000"/>'
        '  <w:left w:val="single" w:sz="4" w:space="0" w:color="000000"/>'
        '  <w:bottom w:val="single" w:sz="4" w:space="0" w:color="000000"/>'
        '  <w:right w:val="single" w:sz="4" w:space="0" w:color="000000"/>'
        '  <w:insideH w:val="single" w:sz="4" w:space="0" w:color="000000"/>'
        '  <w:insideV w:val="single" w:sz="4" w:space="0" w:color="000000"/>'
        '</w:tblBorders>'
    )
    tblPr.append(borders)

    logger.info(
        f"写入表格: {table_info.rows} 行 x {table_info.cols} 列（含合并单元格）"
    )


def build_document(
    tables: list,
    output_path: str,
    page_width_cm: float = 21.0,
    page_height_cm: float = 29.7,
    margin_cm: float = 2.0,
) -> str:
    """将所有表格依次写入一个 DOCX 文件。

    Args:
        tables: list[TableInfo]
        output_path: 输出路径
        page_width_cm: 页面宽度
        page_height_cm: 页面高度
        margin_cm: 页边距

    Returns:
        str: 输出文件路径
    """
    doc = Document()

    section = doc.sections[0]
    section.page_width = Cm(page_width_cm)
    section.page_height = Cm(page_height_cm)
    section.top_margin = Cm(margin_cm)
    section.bottom_margin = Cm(margin_cm)
    section.left_margin = Cm(margin_cm)
    section.right_margin = Cm(margin_cm)

    body_width = page_width_cm - 2 * margin_cm

    for t_idx, table_info in enumerate(tables):
        if table_info.rows == 0 or table_info.cols == 0:
            continue
        if t_idx > 0:
            doc.add_paragraph()
        write_table_to_doc(doc, table_info, body_width)

    doc.save(output_path)
    logger.info(f"文档已保存至: {output_path}")
    return output_path


def build_single_table_document(
    table_info: TableInfo,
    output_path: str,
    page_width_cm: float = 21.0,
    page_height_cm: float = 29.7,
    margin_cm: float = 2.0,
    table_label: str = "",
) -> str:
    """将单个表格写入独立的 DOCX 文件。

    Args:
        table_info: 要写入的表格
        output_path: 输出路径
        page_width_cm: 页面宽度
        page_height_cm: 页面高度
        margin_cm: 页边距
        table_label: 可选标签，写入文档标题

    Returns:
        str: 输出文件路径
    """
    doc = Document()

    section = doc.sections[0]
    section.page_width = Cm(page_width_cm)
    section.page_height = Cm(page_height_cm)
    section.top_margin = Cm(margin_cm)
    section.bottom_margin = Cm(margin_cm)
    section.left_margin = Cm(margin_cm)
    section.right_margin = Cm(margin_cm)

    body_width = page_width_cm - 2 * margin_cm

    if table_label:
        p = doc.add_paragraph()
        run = p.add_run(table_label)
        run.font.size = Pt(12)
        run.font.bold = True

    write_table_to_doc(doc, table_info, body_width)

    doc.save(output_path)
    logger.info(f"单表文档已保存至: {output_path}")
    return output_path