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(