334 lines
12 KiB
Python
334 lines
12 KiB
Python
"""在原始网格里识别表：按空行切块，找表头，合并两级表头，去掉标题、说明和合计行
|
||
|
||
只用规则；判断不了的地方照常产出表，同时记一条警告，调用方据此提示用户
|
||
要看每一行的规则（空行切块、合计行）对整列批量计算；表头、标题、说明只看每块开头或末尾的几行，取成 Python 值判断
|
||
"""
|
||
|
||
import math
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
import pyarrow as pa
|
||
import pyarrow.compute as pc
|
||
|
||
from app.datafiles.base import (
|
||
MAX_TABLE_ROWS,
|
||
Reporter,
|
||
brief,
|
||
no_report,
|
||
rows_exceeded,
|
||
where,
|
||
)
|
||
from app.datafiles.cells import WHITESPACE_CLASS, as_text, classify, mask
|
||
from app.datafiles.grid import Column, Grid
|
||
from app.schemas.data_file import ParseWarning
|
||
|
||
# 只在每块的前几行里找表头
|
||
HEADER_SEARCH_ROWS = 10
|
||
|
||
# 首个或第二个非空单元格是这些词（去掉空白和冒号后）的行是合计行
|
||
TOTAL_LABELS = {
|
||
"合计",
|
||
"总计",
|
||
"小计",
|
||
"共计",
|
||
"总和",
|
||
"平均",
|
||
"平均值",
|
||
"total",
|
||
"subtotal",
|
||
"grandtotal",
|
||
"sum",
|
||
"average",
|
||
}
|
||
# 例如"华东小计""一季度合计"
|
||
TOTAL_SUFFIXES = ("小计", "合计", "总计")
|
||
_TOTAL_SET = pa.array(sorted(TOTAL_LABELS))
|
||
_TOTAL_STRIP = f"[{WHITESPACE_CLASS}:：]"
|
||
_TOTAL_SUFFIX = f"(?:{'|'.join(TOTAL_SUFFIXES)})$"
|
||
# 每个合计字样（含后缀）至少有其中一个字符；英文按转小写之前的大小写都列出
|
||
_TOTAL_HINT = "[计和均tTmMeE]"
|
||
|
||
# 警告里引用的单元格文本最长显示这么多字
|
||
QUOTE_CHARS = 30
|
||
|
||
|
||
@dataclass
|
||
class RawTable:
|
||
sheet: str | None
|
||
# 同一 sheet 切出多张表时的序号，从 1 开始；只有一张表时为空
|
||
part: int | None
|
||
header: list[Any]
|
||
# 数据列，和表头一一对应
|
||
columns: list[Column]
|
||
warnings: list[ParseWarning] = field(default_factory=list)
|
||
|
||
|
||
@dataclass
|
||
class _Piece:
|
||
columns: list[int]
|
||
# 表头所在行在网格中的下标
|
||
header_row: int
|
||
header: list[Any]
|
||
# 数据行在网格中的范围 [start, stop)
|
||
data: list[tuple[int, int]]
|
||
warnings: list[ParseWarning]
|
||
|
||
|
||
def split_tables(
|
||
grid: Grid, report: Reporter = no_report
|
||
) -> tuple[list[RawTable], list[ParseWarning]]:
|
||
"""返回识别出的表和不属于任何表的警告（被丢弃的块）
|
||
|
||
几步都要扫一遍所有行，每步结束报告一次进度，比例来自大表上的实测耗时
|
||
"""
|
||
sheet = grid.sheet
|
||
present = [column.present() for column in grid.columns]
|
||
filled = np.logical_or.reduce(present) if present else np.zeros(grid.height, bool)
|
||
blocks = _blocks(filled)
|
||
used = _used_columns(present, blocks)
|
||
firsts = grid.rows_at(np.array([start for start, _ in blocks], np.int64))
|
||
report(0.3)
|
||
pieces: list[_Piece] = []
|
||
dropped: list[ParseWarning] = []
|
||
for number, ((start, stop), columns, first) in enumerate(
|
||
zip(blocks, used, firsts), start=1
|
||
):
|
||
last = pieces[-1] if pieces else None
|
||
# 表中间夹着空行时，后面的块以数据行开头、列也落在上一张表内，接在上一张表后面
|
||
if (
|
||
last
|
||
and not _is_label_row([first[c] for c in columns])
|
||
and set(columns) <= set(last.columns)
|
||
):
|
||
last.data.append((start, stop))
|
||
continue
|
||
if stop - start < 2:
|
||
dropped.append(
|
||
ParseWarning(
|
||
code="block_dropped",
|
||
message=f"{where(sheet)}第 {start + 1} 行的{_quote(first)}不是表格，已跳过",
|
||
)
|
||
)
|
||
continue
|
||
pieces.append(_read_block(grid, start, stop, columns, sheet))
|
||
report(0.3 + 0.25 * number / len(blocks))
|
||
|
||
tables = []
|
||
for number, piece in enumerate(pieces, start=1):
|
||
table = _finish(grid, piece, sheet)
|
||
report(0.55 + 0.45 * number / len(pieces))
|
||
if table is None:
|
||
dropped.append(
|
||
ParseWarning(
|
||
code="block_dropped",
|
||
message=f"{where(sheet)}第 {piece.header_row + 1} 行开始的表只有表头没有数据，已跳过",
|
||
)
|
||
)
|
||
else:
|
||
tables.append(table)
|
||
report(1.0)
|
||
if len(tables) > 1:
|
||
for part, table in enumerate(tables, start=1):
|
||
table.part = part
|
||
return tables, dropped
|
||
|
||
|
||
def _blocks(filled: np.ndarray) -> list[tuple[int, int]]:
|
||
"""连续的非空行组成的块 [start, stop)"""
|
||
edges = np.flatnonzero(np.diff(np.concatenate(([0], filled, [0])).astype(np.int8)))
|
||
return list(zip(edges[0::2].tolist(), edges[1::2].tolist()))
|
||
|
||
|
||
def _used_columns(
|
||
present: list[np.ndarray], blocks: list[tuple[int, int]]
|
||
) -> list[list[int]]:
|
||
"""每块里有值的列；块之间只隔着空行，按块的起点分段求或即可"""
|
||
if not blocks:
|
||
return []
|
||
starts = np.array([start for start, _ in blocks])
|
||
used = np.array([np.logical_or.reduceat(column, starts) for column in present])
|
||
return [np.flatnonzero(used[:, index]).tolist() for index in range(len(blocks))]
|
||
|
||
|
||
def _read_block(
|
||
grid: Grid, start: int, stop: int, columns: list[int], sheet: str | None
|
||
) -> _Piece:
|
||
# 表头最多在第 HEADER_SEARCH_ROWS 行，两级表头还要看它下面两行
|
||
cells = grid.rows(start, min(stop, start + HEADER_SEARCH_ROWS + 2), columns)
|
||
need = max(1, math.ceil(len(columns) / 2))
|
||
dense = [sum(value is not None for value in row) >= need for row in cells]
|
||
search = range(min(HEADER_SEARCH_ROWS, len(cells)))
|
||
warnings = []
|
||
|
||
# 表头：前几行里第一个够满、并且全是文字的行；标题行通常只有一个单元格，不够满
|
||
header_at = next((i for i in search if dense[i] and _is_label_row(cells[i])), None)
|
||
confident = header_at is not None
|
||
if header_at is None:
|
||
header_at = next((i for i in search if dense[i]), 0)
|
||
warnings.append(
|
||
ParseWarning(
|
||
code="header_uncertain",
|
||
message=f"{where(sheet)}没有找到明确的表头，用第 {start + header_at + 1} 行作表头，可能识别错误",
|
||
)
|
||
)
|
||
if header_at > 0:
|
||
skipped = "".join(
|
||
_quote(cells[i]) for i in range(header_at) if _quote(cells[i])
|
||
)
|
||
warnings.append(
|
||
ParseWarning(
|
||
code="title_skipped",
|
||
message=f"{where(sheet)}{_row_range(start, start + header_at - 1)}是标题，已跳过：{skipped}",
|
||
)
|
||
)
|
||
|
||
header = cells[header_at]
|
||
data_at = header_at + 1
|
||
if (
|
||
confident
|
||
and header_at + 2 < stop - start
|
||
and _is_subheader(cells[header_at], cells[header_at + 1], cells[header_at + 2])
|
||
):
|
||
header = _merge_header(cells[header_at], cells[header_at + 1])
|
||
data_at = header_at + 2
|
||
warnings.append(
|
||
ParseWarning(
|
||
code="header_merged",
|
||
message=f"{where(sheet)}{_row_range(start + header_at, start + header_at + 1)}合并为表头",
|
||
)
|
||
)
|
||
return _Piece(
|
||
columns, start + header_at, header, [(start + data_at, stop)], warnings
|
||
)
|
||
|
||
|
||
def _finish(grid: Grid, piece: _Piece, sheet: str | None) -> RawTable | None:
|
||
ranges = [np.arange(start, stop) for start, stop in piece.data if start < stop]
|
||
rows = np.concatenate(ranges) if ranges else np.zeros(0, np.int64)
|
||
warnings = piece.warnings
|
||
|
||
# 表格下方只有一个文字单元格的行是说明，例如"注：数据来源于……"
|
||
notes = []
|
||
while len(rows) and len(piece.columns) >= 3:
|
||
last = int(rows[-1])
|
||
row = grid.rows(last, last + 1, piece.columns)[0]
|
||
if not _is_note(row):
|
||
break
|
||
notes.append((last, row))
|
||
rows = rows[:-1]
|
||
if notes:
|
||
notes.reverse()
|
||
shown = "".join(_quote(row) for _, row in notes)
|
||
warnings.append(
|
||
ParseWarning(
|
||
code="note_skipped",
|
||
message=f"{where(sheet)}{_row_range(notes[0][0], notes[-1][0])}是说明，已跳过：{shown}",
|
||
)
|
||
)
|
||
|
||
# 合计、小计行不是明细数据，留着会让求和重复计算
|
||
totals = _totals(grid, piece.columns, rows)
|
||
if totals.any():
|
||
listed = brief([str(index + 1) for index in rows[totals].tolist()])
|
||
warnings.append(
|
||
ParseWarning(
|
||
code="total_rows_removed",
|
||
message=f"{where(sheet)}已排除 {int(totals.sum())} 行合计、小计（第 {listed} 行）",
|
||
)
|
||
)
|
||
rows = rows[~totals]
|
||
|
||
if not len(rows):
|
||
return None
|
||
if len(rows) > MAX_TABLE_ROWS:
|
||
raise rows_exceeded(sheet)
|
||
columns = [grid.columns[c].take(rows) for c in piece.columns]
|
||
return RawTable(sheet, None, piece.header, columns, warnings)
|
||
|
||
|
||
def _totals(grid: Grid, columns: list[int], rows: np.ndarray) -> np.ndarray:
|
||
"""rows 里哪些是合计行：每行只看前两个非空单元格，其中有合计字样的文字就是"""
|
||
if not len(rows):
|
||
return np.zeros(0, bool)
|
||
low, high = int(rows[0]), int(rows[-1]) + 1
|
||
seen = np.zeros(high - low, np.int64)
|
||
total = np.zeros(high - low, bool)
|
||
for c in columns:
|
||
column = grid.columns[c].slice(low, high)
|
||
present = column.present()
|
||
open_ = present & (seen < 2)
|
||
if open_.any():
|
||
total |= open_ & column.matches(_is_total_label)
|
||
seen += present
|
||
if (seen >= 2).all():
|
||
break
|
||
return total[rows - low]
|
||
|
||
|
||
def _is_total_label(strings: pa.Array) -> pa.Array:
|
||
"""去掉空白和冒号、转小写后是合计字样；合计字样都含 _TOTAL_HINT 里的字符，先用它筛掉绝大多数值"""
|
||
found = np.zeros(len(strings), bool)
|
||
rows = np.flatnonzero(mask(pc.match_substring_regex(strings, _TOTAL_HINT)))
|
||
if len(rows):
|
||
key = pc.utf8_lower(
|
||
pc.replace_substring_regex(strings.take(rows), _TOTAL_STRIP, "")
|
||
)
|
||
found[rows] = mask(
|
||
pc.or_(
|
||
pc.is_in(key, value_set=_TOTAL_SET),
|
||
pc.match_substring_regex(key, _TOTAL_SUFFIX),
|
||
)
|
||
)
|
||
return pa.array(found)
|
||
|
||
|
||
def _is_label(value: Any) -> bool:
|
||
return isinstance(value, str) and classify(value)[0] == "text"
|
||
|
||
|
||
def _is_label_row(row: list[Any]) -> bool:
|
||
values = [value for value in row if value is not None]
|
||
return bool(values) and all(_is_label(value) for value in values)
|
||
|
||
|
||
def _is_subheader(top: list[Any], sub: list[Any], first_data: list[Any]) -> bool:
|
||
"""第二行也全是文字、紧接着是带数值的数据行，并且两行里有合并单元格留下的空格"""
|
||
return (
|
||
_is_label_row(sub)
|
||
and not _is_label_row(first_data)
|
||
and (None in top or None in sub)
|
||
)
|
||
|
||
|
||
def _merge_header(top: list[Any], sub: list[Any]) -> list[str]:
|
||
"""合并单元格只有左上角有值，第一行横向向前填充后与第二行用 - 连接，例如 销售额-一月"""
|
||
merged, carry = [], None
|
||
for upper, lower in zip(top, sub):
|
||
if upper is not None:
|
||
carry = upper
|
||
parts = [as_text(part) for part in (carry, lower) if part is not None]
|
||
merged.append("-".join(parts))
|
||
return merged
|
||
|
||
|
||
def _is_note(row: list[Any]) -> bool:
|
||
values = [value for value in row if value is not None]
|
||
return len(values) == 1 and _is_label(values[0])
|
||
|
||
|
||
def _quote(row: list[Any]) -> str:
|
||
"""一行里第一个非空单元格，截断后加引号"""
|
||
value = next((value for value in row if value is not None), None)
|
||
if value is None:
|
||
return ""
|
||
text = as_text(value)
|
||
if len(text) > QUOTE_CHARS:
|
||
text = text[:QUOTE_CHARS] + "…"
|
||
return f"「{text}」"
|
||
|
||
|
||
def _row_range(first: int, last: int) -> str:
|
||
return f"第 {first + 1} 行" if first == last else f"第 {first + 1}–{last + 1} 行"
|