arm-statements/arm_statements/tests/builders.py

"""Сборщики тестовых кандидатов (только для тестов пакета)."""

from collections.abc import Mapping, Sequence
from decimal import Decimal

from arm_statements.model import (
    HEADER_FIELDS,
    Cell,
    DocumentType,
    Evidence,
    FormCandidate,
    HeaderCells,
    Source,
    StatementWarning,
)

OcrInput = str | tuple[str, float] | None
ExcelInput = str | Decimal | None
PlaceInput = str | Decimal | tuple[str, float]

_FULL_CONFIDENCE = 100.0
_LATIN_LETTERS = 26


def _text_conf(item: str | tuple[str, float]) -> tuple[str, float]:
    if isinstance(item, tuple):
        return item
    return item, _FULL_CONFIDENCE


def hdr(**fields: str | tuple[str, float]) -> HeaderCells:
    """Шапка из текстов; ref — `hdr:<поле>`."""
    unknown = set(fields) - set(HEADER_FIELDS)
    if unknown:
        raise TypeError(f"неизвестные поля шапки: {sorted(unknown)}")
    cells: dict[str, Cell] = {}
    for name, item in fields.items():
        text, conf = _text_conf(item)
        cells[name] = Cell(text, None, conf, f"hdr:{name}")
    return HeaderCells(**cells)


def _evidence(okud: str | None, title: str | None) -> Evidence:
    return Evidence(
        okud=None if okud is None else Cell(okud, None, _FULL_CONFIDENCE, "evidence:okud"),
        title=None if title is None else Cell(title, None, _FULL_CONFIDENCE, "evidence:title"),
    )


def ocr(  # noqa: PLR0913 - сигнатура задана планом (Task 2)
    rows: Sequence[Sequence[OcrInput]],
    *,
    type_hint: DocumentType | None = None,
    header: HeaderCells | None = None,
    title: str | None = None,
    okud: str | None = None,
    locator: str = "file:t.json;table:0",
) -> FormCandidate:
    """Кандидат OCR: строка r → ячейки с ref `row:<r>`."""
    grid: list[tuple[Cell | None, ...]] = []
    for r, row in enumerate(rows):
        cells: list[Cell | None] = []
        for item in row:
            if item is None:
                cells.append(None)
            else:
                text, conf = _text_conf(item)
                cells.append(Cell(text, None, conf, f"row:{r}"))
        grid.append(tuple(cells))
    return FormCandidate(
        source=Source.OCR,
        locator=locator,
        type_hint=type_hint,
        evidence=_evidence(okud, title),
        header=header if header is not None else HeaderCells(),
        rows=tuple(grid),
    )


def _column_letters(index: int) -> str:
    """0 → A, 25 → Z, 26 → AA."""
    letters = ""
    n = index + 1
    while n > 0:
        n, rem = divmod(n - 1, _LATIN_LETTERS)
        letters = chr(ord("A") + rem) + letters
    return letters


def _sheet_prefix(name: str) -> str:
    """Имя листа в ссылке по правилу Excel: кавычки, если не только буквы/цифры/_ или начинается с цифры."""
    plain = all(ch.isalpha() or ch.isdigit() or ch == "_" for ch in name) and not name[:1].isdigit()
    if plain:
        return name
    return "'" + name.replace("'", "''") + "'"


def _thousands(value: Decimal) -> str:
    return f"{value:,}".replace(",", " ")


def excel(  # noqa: PLR0913 - сигнатура задана планом (Task 2)
    sheets: Sequence[tuple[str, Sequence[Sequence[ExcelInput]]]],
    *,
    type_hint: DocumentType | None = None,
    header: HeaderCells | None = None,
    title: str | None = None,
    okud: str | None = None,
    warnings: Sequence[StatementWarning] = (),
) -> FormCandidate:
    """Кандидат Excel: листы подряд в одной сетке, ref `<лист>!<A1>`."""
    grid: list[tuple[Cell | None, ...]] = []
    breaks: list[int] = []
    for sheet_no, (name, sheet_rows) in enumerate(sheets):
        if sheet_no > 0:
            breaks.append(len(grid))
        prefix = _sheet_prefix(name)
        for r, row in enumerate(sheet_rows):
            cells: list[Cell | None] = []
            for c, item in enumerate(row):
                ref = f"{prefix}!{_column_letters(c)}{r + 1}"
                if item is None:
                    cells.append(None)
                elif isinstance(item, Decimal):
                    cells.append(Cell(_thousands(item), item, _FULL_CONFIDENCE, ref))
                else:
                    cells.append(Cell(item, None, _FULL_CONFIDENCE, ref))
            grid.append(tuple(cells))
    return FormCandidate(
        source=Source.EXCEL,
        locator="sheet:" + ",".join(name for name, _ in sheets),
        type_hint=type_hint,
        evidence=_evidence(okud, title),
        header=header if header is not None else HeaderCells(),
        rows=tuple(grid),
        segment_breaks=tuple(breaks),
        warnings=tuple(warnings),
    )


def place(cells: Mapping[int, PlaceInput], width: int | None = None) -> list[PlaceInput | None]:
    """Строка с ячейками в заданных столбцах, остальные None."""
    size = width if width is not None else (max(cells) + 1 if cells else 0)
    return [cells.get(i) for i in range(size)]