diff --git a/api/src/core/PDFWorker/utils.py b/api/src/core/PDFWorker/utils.py index 0d42047..c39bbe9 100644 --- a/api/src/core/PDFWorker/utils.py +++ b/api/src/core/PDFWorker/utils.py @@ -1,11 +1,14 @@ from __future__ import annotations +import re import uuid from dataclasses import dataclass, asdict from datetime import date, datetime from decimal import Decimal, InvalidOperation from typing import Optional, Literal, TypedDict +from loguru import logger + from src.schemas.enums import StatementType, FslRowKind from .consts import codes_mapping @@ -52,6 +55,27 @@ class ExtractResult(TypedDict): skipped_files: list[tuple[str, dict]] +@dataclass(frozen=True) +class ColumnLayout: + """Индексы значимых столбцов таблицы OCR, определённые по строке заголовка.""" + label_idx: Optional[int] + code_idx: int + value_idxs: tuple[int, ...] + + +@dataclass +class ParsedRow: + row: dict + cells: dict[int, dict] + layout: ColumnLayout + kind: FslRowKind + + +# column_index в БД: 0 - наименование, 1 - код, значения начинаются с 2 +_FIRST_VALUE_COLUMN = 2 +_YEAR_RE = re.compile(r"\d{4}") + + _TABLE_TYPE_MAP: dict[str, StatementType] = { "Бухгалтерский баланс": StatementType.BALANCE, "Отчет о финансовых результатах": StatementType.INCOME, @@ -106,16 +130,61 @@ def _parse_amount(raw: str) -> tuple[Optional[Decimal], bool]: return (None, is_verified) -def _classify_row(cells: list[dict]) -> FslRowKind | None: - if not (str(cells[1]['value']).strip()): +def _normalize_header(value) -> str: + text = str(value or '').lower() + text = ''.join(c if c.isalnum() else ' ' for c in text) + return ' '.join(text.split()) + + +def _is_code_header(text: str) -> bool: + return text == 'код' or text.startswith('код ') + + +def _is_value_header(text: str) -> bool: + return bool(_YEAR_RE.search(text)) or text.startswith(('на ', 'за ')) + + +def _cells_by_index(row: dict) -> dict[int, dict]: + return {cell['column_index']: cell for cell in row['cells']} + + +def _cell_value(cells: dict[int, dict], idx: Optional[int]) -> str: + if idx is None or idx not in cells: + return '' + return str(cells[idx]['value'] or '') + + +def _build_layout(cells: dict[int, dict]) -> Optional[ColumnLayout]: + """Возвращает раскладку столбцов, если строка является заголовком (есть ячейка «Код»).""" + headers = {idx: _normalize_header(cells[idx]['value']) for idx in sorted(cells)} + code_idx = next((idx for idx, text in headers.items() if _is_code_header(text)), None) + if code_idx is None: + return None + + label_idx = next((idx for idx, text in headers.items() if 'наименование' in text), None) + if label_idx is None: + left_filled = [idx for idx, text in headers.items() if idx < code_idx and text] + label_idx = left_filled[-1] if left_filled else None + + value_idxs = tuple( + idx for idx, text in headers.items() + if idx > code_idx and _is_value_header(text) + ) + return ColumnLayout(label_idx=label_idx, code_idx=code_idx, value_idxs=value_idxs) + + +def _classify_row(cells: dict[int, dict], layout: ColumnLayout) -> FslRowKind | None: + code = _cell_value(cells, layout.code_idx).strip() + if not code: return None try: - int(cells[1]["value"]) + int(code) except (TypeError, ValueError): return FslRowKind.HEADER - if cells[0]["value"] == "Баланс": + label = _cell_value(cells, layout.label_idx) + if label == "Баланс": return FslRowKind.BALANCE - elif cells[0]["value"].lower().startswith("итого"): + elif label.lower().startswith("итого"): return FslRowKind.TOTAL return FslRowKind.ROW @@ -129,19 +198,46 @@ def _keep_digits_only(raw: str) -> str: return res.strip() -def _resolve_table_type(rows) -> StatementType: +def _parse_table_rows(table_rows: list[dict]) -> Optional[list[ParsedRow]]: + """Разбирает строки таблицы по раскладке из заголовка. + + Раскладка берётся из строки с ячейкой «Код»; каждый повторный заголовок + перезаписывает индексы столбцов. В результат попадает только первый заголовок, + строки до первого заголовка игнорируются. Если заголовка нет, возвращает None. + """ + layout: Optional[ColumnLayout] = None + parsed_rows: list[ParsedRow] = [] + + for row in table_rows: + cells = _cells_by_index(row) + header_layout = _build_layout(cells) + if header_layout is not None: + if layout is None: + parsed_rows.append(ParsedRow(row, cells, header_layout, FslRowKind.HEADER)) + layout = header_layout + continue + if layout is None: + continue + + kind = _classify_row(cells, layout) + if kind is None or kind == FslRowKind.HEADER: + continue + parsed_rows.append(ParsedRow(row, cells, layout, kind)) + + return parsed_rows if layout is not None else None + + +def _resolve_table_type(parsed_rows: list[ParsedRow]) -> StatementType: cell_codes = { '1': 0, '2': 0 } try: - for r in rows: - cells = list(sorted(r["cells"], key=lambda x: x['column_index'])) - kind = _classify_row(cells) - if kind is None or kind == FslRowKind.HEADER: + for parsed in parsed_rows: + if parsed.kind == FslRowKind.HEADER: continue else: - first_symbol = str(cells[1]['value'])[0] + first_symbol = _cell_value(parsed.cells, parsed.layout.code_idx).strip()[0] cell_codes.setdefault(first_symbol, 0) cell_codes[first_symbol] += 1 if cell_codes['1'] > cell_codes['2']: @@ -196,36 +292,34 @@ def _process_file( rows: list[DbRow] = [] values: list[DbValue] = [] - for table in data["table_structure"]: + for table_number, table in enumerate(data["table_structure"]): + parsed_rows = _parse_table_rows(table["rows"]) + if parsed_rows is None: + logger.warning( + "В таблице не найден заголовок со столбцом «Код», таблица пропущена. " + "file_uid={file_uid}, table_number={table_number}, table_label={table_label}", + file_uid=file.get("file_uid"), + table_number=table_number, + table_label=(table.get("label") or {}).get("form_name"), + ) + continue + meta = Meta( id=uuid.uuid4(), - statement_type=_resolve_table_type(table["rows"]), + statement_type=_resolve_table_type(parsed_rows), **shared, ) metas.append(asdict(meta)) - header_found = False - for row in table["rows"]: - cells = list(sorted(row["cells"], key=lambda x: x['column_index'])) - if len(cells) < 3: - raise ValueError( - f"Row {row['row_index']} in file {file['file_uid']} " - f"has {len(cells)} cells (< 3)." - ) - - kind = _classify_row(cells) - if kind is None or (kind == FslRowKind.HEADER and header_found): - continue - elif (kind == FslRowKind.HEADER and not(header_found)): - header_found = True - label_cell = cells[0] - code_cell = cells[1] + for parsed in parsed_rows: + row, cells, layout, kind = parsed.row, parsed.cells, parsed.layout, parsed.kind + code_cell = cells[layout.code_idx] is_line_code_valid = True line_code = _clear_value(code_cell["value"]) line_name = codes_mapping.get(line_code, '') if line_name == '': is_line_code_valid = False - line_name = label_cell['value'] + line_name = _cell_value(cells, layout.label_idx) line_code_recognized = _is_recognized(code_cell["confidence"]) and _clear_value(code_cell['value']) == line_code and is_line_code_valid @@ -245,15 +339,19 @@ def _process_file( rows.append(asdict(db_row)) - for cell in cells: - if cell["column_index"] < 2: - continue - raw = cell["value"] + for value_number, source_idx in enumerate(layout.value_idxs): + cell = cells.get(source_idx) parsed_amount = None - if kind == FslRowKind.HEADER: + if cell is None: + # ячейки под столбцом заголовка нет — пишем пустое значение, чтобы не сдвигать столбцы + raw = '' + is_recognized = False + elif kind == FslRowKind.HEADER: + raw = cell["value"] is_recognized = (_clear_value(raw) == _remove_spec_space(raw) and _is_recognized(cell["confidence"])) raw = _clear_value(raw) else: + raw = cell["value"] parsed_amount, is_num_recognized = _parse_amount(raw) is_recognized = _is_recognized(cell["confidence"]) and is_num_recognized raw = _clear_value(raw) @@ -261,7 +359,7 @@ def _process_file( line_id=db_row.id, amount=parsed_amount, string_value=raw, - column_index=cell["column_index"], + column_index=_FIRST_VALUE_COLUMN + value_number, is_recognized=is_recognized, ) values.append(