pylib_ai_arm_statements/arm_statements/tests/test_layout.py

from decimal import Decimal as D

import pytest

from arm_statements.layout import (
    code_column,
    find_code_column,
    find_header_rows,
    find_name_column,
    first_code_row,
    is_numbering_row,
    name_cell,
    segment_roles,
)
from arm_statements.loader import CodeSection, Rules
from arm_statements.model import Cell, DocumentType, RowRole
from arm_statements.tests.builders import excel, ocr, place

BALANCE_ROWS = [
    ["Наименование показателя", "Код", "На 30.09.2025", "На 31.12.2024"],
    ["1", "2", "3", "4"],
    ["Нематериальные активы", "1110", "10", "9"],
    ["Основные средства", "1150", "20", "19"],
    ["Итого по разделу I", "1100", "30", "28"],
]


def test_code_column_by_caption(rules: Rules) -> None:
    cand = ocr(BALANCE_ROWS)
    assert (find_code_column(cand.rows, 0, 4, rules), first_code_row(cand.rows, 0, 4, 1, rules)) == (1, 2)


def test_title_block_kody_is_not_code_column(rules: Rules) -> None:
    title = [
        place({5: v}) for v in ["Коды", "0710001", "30 09 2025", "12345678", "7700000000", "12267", "16", "384", "", ""]
    ]
    table = [
        place({0: n, 2: c, 3: D(a)})
        for n, c, a in [
            ("Запасы", "1210", "5"),
            ("НДС", "1220", "1"),
            ("Дебиторы", "1230", "4"),
            ("Итого", "1200", "10"),
        ]
    ]
    cand = excel([("Баланс", [*title, *table])])
    assert [find_code_column(cand.rows, 0, len(cand.rows) - 1, rules)] == [2]


def test_two_line_document(rules: Rules) -> None:
    cand = ocr([["Наименование", "Код", "2025"], ["Выручка", "2110", "5"], ["Прибыль", "2400", "1"]])
    assert find_code_column(cand.rows, 0, 2, rules) == 1


def test_relaxed_search_finds_single_code(rules: Rules) -> None:
    cand = ocr([["Чистая прибыль", "", "2400", "7", "6"]])
    assert find_code_column(cand.rows, 0, 0, rules) is None
    assert [find_code_column(cand.rows, 0, 0, rules, relaxed_section=CodeSection.INCOME)] == [2]


def test_public_code_column(rules: Rules) -> None:
    assert code_column(ocr(BALANCE_ROWS).rows, rules) == 1
    single = ocr([["Чистая прибыль", "", "2400", "7", "6"]]).rows
    assert (code_column(single, rules), code_column(single, rules, relaxed=DocumentType.INCOME)) == (None, 2)
    assert code_column(single, rules, relaxed=DocumentType.BALANCE) is None and code_column((), rules) is None
    with pytest.raises(ValueError, match="DISCLOSURE"):
        code_column(single, rules, relaxed=DocumentType.DISCLOSURE)


def test_numbering_rows(rules: Rules) -> None:
    def row(*texts: str) -> tuple[Cell | None, ...]:
        return ocr([list(texts)]).rows[0]

    assert is_numbering_row(row("1", "2", "3", "4", "5"), rules)
    assert is_numbering_row(row("1", "2", "З", "4"), rules)
    assert is_numbering_row(row("2", "3", "4", "5"), rules)
    assert is_numbering_row(row("1", "2", "4", "5"), rules)
    assert not is_numbering_row(row("1", "3", "5"), rules)
    assert not is_numbering_row(row("1", "2"), rules)


def test_header_row_roles_and_title(rules: Rules) -> None:
    rows = [["Бухгалтерский баланс", "", "", ""], *BALANCE_ROWS]
    cand = ocr(rows)
    header = find_header_rows(cand.rows, 0, 3, 1, rules)
    assert header == (1,)
    roles = segment_roles(cand.rows, 0, 5, header, rules)
    assert roles == [RowRole.TITLE, RowRole.HEADER, RowRole.NUMBERING, RowRole.DATA, RowRole.DATA, RowRole.DATA]
    assert RowRole.TITLE not in segment_roles(cand.rows, 3, 5, (), rules)


def test_two_row_header_is_glued(rules: Rules) -> None:
    cand = ocr(
        [
            ["", "", "На 30 сентября", "На 31 декабря"],
            ["Наименование", "Код", "2025 г.", "2024 г."],
            ["Запасы", "1210", "5", "4"],
        ]
    )
    assert find_header_rows(cand.rows, 0, 2, 1, rules) == (0, 1)
    single = ocr(
        [
            ["Раздел", "", "", ""],
            ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024"],
            ["Запасы", "1210", "5", "4"],
        ]
    )
    assert find_header_rows(single.rows, 0, 2, 1, rules) == (1,)


def test_name_column_without_caption(rules: Rules) -> None:
    cand = ocr(
        [
            ["5.1", "Запасы", "1210", "5"],
            ["5.2", "Дебиторская задолженность", "1230", "4"],
            ["", "Итого", "1200", "9"],
        ]
    )
    assert find_name_column(cand.rows, (), 0, 2, 2, rules) == (1, 0)


def test_header_without_code_caption_skips_numbering(rules: Rules) -> None:
    cand = ocr(
        [
            ["Титул 2025", "", "", ""],
            ["Показатель", "", "На 30.09.2025", "На 31.12.2024"],
            ["1", "2", "3", "4"],
            ["Запасы", "1210", "5", "4"],
        ]
    )
    assert find_header_rows(cand.rows, 0, 3, 1, rules) == (1,)
    assert find_header_rows(cand.rows, 3, 3, 1, rules) == ()


def test_name_cell_falls_back_to_left_text(rules: Rules) -> None:
    row = ocr([["5.1", "Запасы", "", "1210"]]).rows[0]
    cell = name_cell(row, 2, 3, 0)
    assert cell is not None and cell.text == "Запасы"
    assert name_cell(ocr([["5.1", "", "", "1210"]]).rows[0], 2, 3, 0) is None