proposa-ai/pdf_table_to_docx/table_parser.py

"""
table_parser.py - PDF 表格解析模块

负责：
- 使用 pdfplumber 从页面中提取表格
- 根据单元格坐标信息检测合并单元格（跨行/跨列）
- 提取每个单元格的文本内容
"""

import logging
from dataclasses import dataclass, field
from typing import Optional

import pdfplumber.table

logger = logging.getLogger(__name__)


@dataclass
class CellInfo:
    """单个表格单元格的信息"""
    text: str                       # 单元格文本内容
    row_span: int = 1               # 跨行数（垂直合并）
    col_span: int = 1               # 跨列数（水平合并）
    bbox: Optional[tuple] = None    # 原始坐标 (x0, top, x1, bottom)


@dataclass
class TableInfo:
    """解析后的表格结构"""
    rows: int = 0                           # 总行数
    cols: int = 0                           # 总列数
    cells: list = field(default_factory=list)  # 二维列表: cells[row][col] -> CellInfo
    col_widths: list = field(default_factory=list)  # 每列宽度（相对比例）
    bbox: Optional[tuple] = None            # 表格在页面上的坐标
    page_num: int = 0                       # 所在页码（0-indexed）
    cross_page: bool = False                # 是否跨页表格
    original_pages: list = field(default_factory=list)  # 跨页时记录所有页码（1-indexed）


BBOX_EPSILON = 5.0  # pt，bbox 判断容差（约 1.8mm）


def _detect_merge_spans(grid, n_rows, n_cols, rows_bottoms=None):
    """检测合并单元格的跨行跨列范围，自动处理 T 形冲突。

    检测策略（基于 pdfplumber 的网格特性）：
    - pdfplumber 将表格表示为网格 grid[r][c]
    - 一个合并单元格在它的起始位置 (r, c) 存储 bbox（非 None）
    - 被它覆盖的其他网格位置均为 None
    - 从起始位置向右/向下扫描连续的 None 确定 col_span 和 row_span

    BBox 验证（核心修复）：
    - 仅靠 None 扫描可能产生假阳性纵向合并（如 page 24 的"合计"行，
      其中下方 None 单元格实际属于水平合并而非纵向合并）。
    - 因此，对 row_span > 1 的候选，验证单元格 bbox 底部确实延伸到
      目标行的底部，否则缩减行跨度。

    T 形冲突处理：
    - DOCX 不允许非矩形合并。
    - 按 row-major 顺序认领区域，后处理的合并若冲突则自动缩减。

    Args:
        grid: grid[r][c] = bbox tuple or None
        n_rows: 总行数
        n_cols: 总列数
        rows_bottoms: list[float] 每行的底部 y 坐标，用于 bbox 验证

    Returns:
        dict: {(r, c): (row_span, col_span)} 映射
    """
    # ---- 第一遍：初步检测所有可能的合并 ----
    candidates = []  # [(r, c, row_span, col_span), ...]
    for r in range(n_rows):
        for c in range(n_cols):
            cell = grid[r][c]
            if cell is None:
                continue

            x0, top, x1, bottom = cell

            # 列跨度：向右扫描连续 None
            col_span = 1
            for cc in range(c + 1, n_cols):
                if grid[r][cc] is None:
                    col_span += 1
                else:
                    break

            # 行跨度：向下扫描连续 None
            row_span = 1
            for rr in range(r + 1, n_rows):
                if grid[rr][c] is None:
                    row_span += 1
                else:
                    break

            # ---- BBox 验证 ----
            # 纵向合并验证：bbox 底部必须延伸到目标行的底部
            if row_span > 1 and rows_bottoms is not None:
                target_bottom = rows_bottoms[r + row_span - 1]
                if bottom < target_bottom - BBOX_EPSILON:
                    # bbox 不够长，往回缩减
                    reduced = False
                    for rs in range(row_span - 1, 0, -1):
                        if bottom >= rows_bottoms[r + rs - 1] - BBOX_EPSILON:
                            row_span = rs
                            reduced = True
                            break
                    if not reduced:
                        row_span = 1

            # 验证矩形完整性
            if row_span > 1 or col_span > 1:
                all_clear = True
                for rr in range(r, r + row_span):
                    for cc in range(c, c + col_span):
                        if rr == r and cc == c:
                            continue
                        if grid[rr][cc] is not None:
                            all_clear = False
                            break
                    if not all_clear:
                        break
                if not all_clear:
                    row_span = 1
                    col_span = 1

            candidates.append((r, c, row_span, col_span))

    # ---- 第二遍：按 row-major 排序，解决冲突 ----
    candidates.sort(key=lambda x: (x[0], x[1]))

    claimed = set()
    merge_spans = {}

    for r, c, row_span, col_span in candidates:
        if (r, c) in claimed:
            merge_spans[(r, c)] = (0, 0)
            continue

        if row_span == 1 and col_span == 1:
            claimed.add((r, c))
            merge_spans[(r, c)] = (1, 1)
            continue

        actual_row_span = row_span
        actual_col_span = col_span

        # 先缩减列跨度：检查每行中是否有已被认领的列
        for rr in range(r, r + actual_row_span):
            for cc in range(c + 1, c + actual_col_span):
                if (rr, cc) in claimed:
                    actual_col_span = cc - c
                    break
            if rr > r:
                for cc in range(c + 1, c + actual_col_span):
                    if (rr, cc) in claimed:
                        actual_col_span = min(actual_col_span, cc - c)
                        break

        # 再缩减行跨度：检查每列中是否有已被认领的行
        for cc in range(c, c + actual_col_span):
            for rr in range(r + 1, r + actual_row_span):
                if (rr, cc) in claimed:
                    actual_row_span = rr - r
                    break

        actual_row_span = max(1, actual_row_span)
        actual_col_span = max(1, actual_col_span)

        # 认领区域
        for rr in range(r, r + actual_row_span):
            for cc in range(c, c + actual_col_span):
                claimed.add((rr, cc))

        merge_spans[(r, c)] = (actual_row_span, actual_col_span)

    return merge_spans


def parse_table(
    table: pdfplumber.table.Table,
    page_num: int = 0,
) -> TableInfo:
    """将一个 pdfplumber Table 对象解析为 TableInfo，检测合并单元格。

    pdfplumber 的 Table 对象中，被合并"覆盖"的网格位置返回 None，
    我们利用这个信息来推断每个单元格的 row_span 和 col_span。

    核心检测策略：
    1. 遍历表格网格 grid[r][c]。
    2. 对于非 None 的单元格，向右扫描连续的 None 来计算列跨度，
       向下扫描连续的 None 来计算行跨度。
    3. 验证合并矩形区域的完整性。
    """
    n_rows = len(table.rows)
    n_cols = len(table.columns)

    if n_rows == 0 or n_cols == 0:
        logger.warning(f"第 {page_num+1} 页发现空表格，跳过")
        return TableInfo()

    # 构建原始网格：grid[r][c] = bbox tuple or None
    grid: list = []
    for r in range(n_rows):
        row_cells = []
        for c in range(n_cols):
            cell = table.rows[r].cells[c]
            row_cells.append(cell)
        grid.append(row_cells)

    # ---- 1. 检测合并范围 ----
    rows_bottoms = [table.rows[r].bbox[3] for r in range(n_rows)]
    merge_spans = _detect_merge_spans(
        grid, n_rows, n_cols,
        rows_bottoms=rows_bottoms,
    )

    # ---- 2. 提取文字内容 ----
    try:
        raw_text = table.extract()
    except Exception as e:
        logger.warning(f"表格文字提取失败: {e}")
        raw_text = None

    # ---- 2.5 清理不必要的换行符（PDF 换行包装 vs 结构性换行） ----
    if raw_text:
        try:
            from .text_scorer import get_scorer
            _scorer = get_scorer()
            _score_fn = _scorer.score
        except Exception:
            _score_fn = None

        def _clean_nl(text: str) -> str:
            """去除 PDF 换行包装引入的 \\n，保留结构性换行。"""
            if '\n' not in text:
                return text
            parts = text.split('\n')

            # 规则1：某行只有 1 个字符 → 包装换行，去
            if any(len(p.strip()) <= 1 for p in parts):
                return text.replace('\n', '')

            # 规则2：括号分裂如 "（万\n元）" → 去
            open_br = set('（({[〈《「『【')
            close_br = set('）)}]〉》」』】')
            for i in range(1, len(parts)):
                prev, curr = parts[i-1].strip(), parts[i].strip()
                if prev and curr:
                    if prev[-1] in open_br or curr[0] in close_br:
                        return text.replace('\n', '')

            # 规则3：分值模式 "文字\n数字" → 保
            last = parts[-1].strip()
            if last.replace('.','').isdigit() and len(parts) > 1:
                before = '\n'.join(parts[:-1]).strip()
                if len(before) > 1:
                    return text

            # 规则4：NLP 评分（文本够长时）
            cleaned = text.replace('\n', '')
            if _score_fn and len(cleaned) >= 6:
                try:
                    s_c = _score_fn(cleaned)
                    line_scores = [_score_fn(p.strip()) for p in parts if p.strip()]
                    if line_scores:
                        max_l = max(line_scores)
                        if s_c > max_l * 1.2:
                            return cleaned
                        if s_c > 0.01 and max_l < 0.008:
                            return cleaned
                except Exception:
                    pass
            return text  # 保守保留

        for r in range(n_rows):
            for c in range(n_cols):
                raw = raw_text[r][c]
                if raw:
                    raw_text[r][c] = _clean_nl(raw)

    # ---- 3. 构建 CellInfo 网格 ----
    cells: list = []
    for r in range(n_rows):
        row_cells = []
        for c in range(n_cols):
            if grid[r][c] is None:
                # 被合并覆盖，填充占位
                row_cells.append(CellInfo(text="", row_span=0, col_span=0))
            else:
                rs, cs = merge_spans.get((r, c), (1, 1))

                # 从 extract() 结果中获取文本
                text = ""
                if raw_text and r < len(raw_text) and c < len(raw_text[r]):
                    raw = raw_text[r][c]
                    if raw is not None:
                        text = raw

                row_cells.append(CellInfo(
                    text=text,
                    row_span=rs,
                    col_span=cs,
                    bbox=grid[r][c],
                ))
        cells.append(row_cells)

    # ---- 4. 计算每列相对宽度 ----
    col_widths = []
    for c in range(n_cols):
        col_bbox = table.columns[c].bbox  # (x0, top, x1, bottom)
        w = col_bbox[2] - col_bbox[0]
        col_widths.append(w)

    total_w = sum(col_widths)
    if total_w > 0:
        col_widths = [w / total_w for w in col_widths]

    return TableInfo(
        rows=n_rows,
        cols=n_cols,
        cells=cells,
        col_widths=col_widths,
        bbox=table.bbox,
        page_num=page_num,
    )


def tables_on_page(page, page_num: int) -> list:
    """提取一页中的所有表格。

    Returns:
        list[TableInfo]: 页面上所有解析后的表格列表
    """
    try:
        raw_tables = page.find_tables()
    except Exception as e:
        logger.error(f"第 {page_num+1} 页表格检测失败: {e}")
        return []

    results = []
    for table in raw_tables:
        try:
            info = parse_table(table, page_num)
            if info.rows > 0 and info.cols > 0:
                results.append(info)
        except Exception as e:
            logger.warning(f"第 {page_num+1} 页中某个表格解析失败: {e}")

    return results