arm-statements/arm_statements/codes.py

"""Разбор кода строки и отрезки кодов форм."""

import re
from collections import Counter
from collections.abc import Iterable, Mapping, Sequence
from typing import NamedTuple

from arm_statements.loader import CodeSection, Rules
from arm_statements.model import DocumentType
from arm_statements.text import replace_lookalikes

_CODE = re.compile(r"\d{4,5}(?:\.0)?")
_BAR = "|"


class CodeParse(NamedTuple):
    code: str | None
    fixed: bool  # при разборе заменён двойник


def code_parts(text: str | None) -> tuple[str, ...]:
    """Непустые строки ячейки, обрезанные."""
    if text is None:
        return ()
    return tuple(part.strip() for part in text.splitlines() if part.strip())


def section_of(code: str, rules: Rules) -> CodeSection | None:
    head = int(code[:4])
    for item in rules.code_ranges.ranges:
        if item.low <= head <= item.high:
            return item.section
    return None


def _strip_edges(text: str) -> str:
    start, end = 0, len(text)
    while start < end and not text[start].isalnum():
        start += 1
    while end > start and not text[end - 1].isalnum():
        end -= 1
    return text[start:end]


def _as_code(text: str, rules: Rules) -> str | None:
    core = _strip_edges(text)
    if _CODE.fullmatch(core) is None:
        return None
    code = core.split(".", 1)[0]
    return code if section_of(code, rules) is not None else None


def _parse_line(line: str, rules: Rules) -> CodeParse:
    compact = "".join(line.split())
    # сначала `|` как двойник единицы; если результата-кода нет — `|` как краевая пунктуация
    for candidate in (compact, compact.replace(_BAR, "")):
        replaced = replace_lookalikes(candidate, rules.lookalikes)
        code = _as_code(replaced, rules)
        if code is not None:
            return CodeParse(code, replaced != candidate)
    return CodeParse(None, False)


def parse_code(text: str | None, rules: Rules) -> CodeParse:
    """Однострочный текст → код; многострочный → `CodeParse(None, False)`."""
    parts = code_parts(text)
    if len(parts) != 1:
        return CodeParse(None, False)
    return _parse_line(parts[0], rules)


def is_code_cell(text: str | None, rules: Rules) -> bool:
    parts = code_parts(text)
    return bool(parts) and all(_parse_line(part, rules).code is not None for part in parts)


def majority_section(codes: Iterable[str], rules: Rules) -> CodeSection | None:
    """Отрезок большинства; ничья — первый по порядку отрезков словаря."""
    counts: Counter[CodeSection] = Counter()
    for code in codes:
        section = section_of(code, rules)
        if section is not None:
            counts[section] += 1
    best: CodeSection | None = None
    for item in rules.code_ranges.ranges:
        if counts[item.section] and (best is None or counts[item.section] > counts[best]):
            best = item.section
    return best


def document_type_of(section: CodeSection) -> DocumentType:
    mapping: Mapping[CodeSection, DocumentType] = {
        CodeSection.BALANCE: DocumentType.BALANCE,
        CodeSection.INCOME: DocumentType.INCOME,
    }
    return mapping.get(section, DocumentType.DISCLOSURE)


def sheet_joins(
    form_type: DocumentType,
    sheet_codes: Sequence[str],
    sheet_amount_columns: int,
    form_amount_columns: int,
    rules: Rules,
) -> bool:
    """Лист без маркера и шапки присоединяется к форме: коды в её отрезке; при ≥3 кодах колонки сумм совпадают."""
    if len(sheet_codes) < rules.code_ranges.min_codes_continuation:
        return False
    form_section = {DocumentType.BALANCE: CodeSection.BALANCE, DocumentType.INCOME: CodeSection.INCOME}.get(form_type)
    for code in sheet_codes:
        section = section_of(code, rules)
        if section is None or document_type_of(section) is not form_type:
            return False
        if form_section is not None and section is not form_section:
            return False
    if len(sheet_codes) >= rules.code_ranges.min_codes:
        return sheet_amount_columns == form_amount_columns
    return True