arm-scan-gateway/api/src/normalizer/rules.py

"""Справочники нормализатора OCR (D-7, §11.1 п. 1, 6): yaml в `SCAN_RULES_DIR`, загрузка один раз при старте.

`statement_lines.yaml` — редакции форм (`66n` до 2024-12-31, `fsbu4` с 2025-01-01) с таблицей код → наименование;
`ocr_lookalikes.yaml` — двойники букв и цифр; `statement_labels.yaml` — подстроки `label` таблицы → тип документа.
Битый файл — `ValueError` с именем файла (сервис не стартует); повторяющиеся ключи yaml (код строки дважды) — тоже
ошибка (№33).
"""

from collections.abc import Iterable
from datetime import date
from pathlib import Path
from typing import Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator, model_validator

__all__ = ["Edition", "OcrRules", "load_rules"]

LINES_FILE = "statement_lines.yaml"
LOOKALIKES_FILE = "ocr_lookalikes.yaml"
LABELS_FILE = "statement_labels.yaml"
DocumentKind = Literal["BALANCE", "INCOME"]


class Edition(BaseModel):
    """Редакция форм: интервал дат отчёта (включительно, открытый край — `None`) и таблица код → наименование."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    id: str = Field(min_length=1)
    valid_from: date | None = None
    valid_until: date | None = None
    lines: dict[str, str] = Field(default_factory=dict)

    @field_validator("lines")
    @classmethod
    def _check_lines(cls, lines: dict[str, str]) -> dict[str, str]:
        for code, name in lines.items():
            if not (code.isdigit() and 4 <= len(code) <= 5):  # noqa: PLR2004 — коды строк форм: 4–5 цифр
                raise ValueError(f"код строки {code!r}: ожидаются 4–5 цифр")
            if not name.strip():
                raise ValueError(f"код строки {code!r}: пустое наименование")
        return lines

    def covers(self, report_date: date) -> bool:
        return (self.valid_from is None or self.valid_from <= report_date) and (
            self.valid_until is None or report_date <= self.valid_until
        )


class _LinesFile(BaseModel):
    model_config = ConfigDict(extra="forbid")
    editions: list[Edition] = Field(min_length=1)

    @model_validator(mode="after")
    def _unique_ids(self) -> "_LinesFile":
        ids = [edition.id for edition in self.editions]
        if len(ids) != len(set(ids)):
            raise ValueError(f"id редакций повторяются: {ids}")
        return self


class _LookalikesFile(BaseModel):
    model_config = ConfigDict(extra="forbid")
    digits: dict[str, str]

    @field_validator("digits")
    @classmethod
    def _single_chars(cls, digits: dict[str, str]) -> dict[str, str]:
        for char, digit in digits.items():
            if len(char) != 1 or len(digit) != 1 or not digit.isdigit():
                raise ValueError(f"двойник {char!r} -> {digit!r}: ожидается один символ и одна цифра")
        return digits


class _LabelsFile(BaseModel):
    model_config = ConfigDict(extra="forbid")
    labels: dict[str, DocumentKind]


class OcrRules(BaseModel):
    """Все справочники нормализатора; передаётся в конвейер параметром."""

    model_config = ConfigDict(frozen=True)

    editions: list[Edition]
    lookalikes: dict[str, str]
    labels: dict[str, DocumentKind]

    def edition_for(self, report_date: date | None) -> Edition | None:
        """Редакция по дате отчёта (первая, чей интервал покрывает дату); `None` без даты или вне интервалов."""
        if report_date is None:
            return None
        return next((edition for edition in self.editions if edition.covers(report_date)), None)

    def edition_by_codes(self, codes: Iterable[str]) -> Edition | None:
        """Редакция с наибольшим числом известных кодов; при равенстве — первая; `None`, если совпадений нет."""
        wanted = set(codes)
        best: Edition | None = None
        best_hits = 0
        for edition in self.editions:
            hits = len(wanted & edition.lines.keys())
            if hits > best_hits:
                best, best_hits = edition, hits
        return best

    def known_codes(self) -> frozenset[str]:
        return frozenset(code for edition in self.editions for code in edition.lines)

    def editions_with(self, code: str) -> list[str]:
        return [edition.id for edition in self.editions if code in edition.lines]


class _UniqueKeyLoader(yaml.SafeLoader):
    """`SafeLoader`, отклоняющий повторяющиеся ключи отображений (№33): `safe_load` молча оставляет последнее значение,
    и повторный код в справочнике терялся бы без следа."""

    def construct_mapping(self, node: yaml.MappingNode, deep: bool = False) -> dict[object, object]:
        seen: set[object] = set()
        for key_node, _ in node.value:
            key = self.construct_object(key_node, deep=True)
            try:
                duplicate = key in seen
                seen.add(key)
            except TypeError:  # нехешируемый ключ — его отклонит базовый `construct_mapping`
                continue
            if duplicate:
                raise yaml.constructor.ConstructorError(
                    "при разборе отображения", node.start_mark, f"повторяющийся ключ {key!r}", key_node.start_mark
                )
        return super().construct_mapping(node, deep=deep)


def _read_yaml(directory: Path, name: str) -> object:
    path = directory / name
    try:
        text = path.read_text(encoding="utf-8")
    except OSError as exc:
        raise ValueError(f"справочник {name} не прочитан ({path}): {exc}") from exc
    try:
        return yaml.load(text, Loader=_UniqueKeyLoader)  # SafeLoader с проверкой дублей
    except yaml.YAMLError as exc:
        raise ValueError(f"справочник {name} не разобран как YAML: {exc}") from exc


def load_rules(directory: Path | str) -> OcrRules:
    """Загрузить три справочника из каталога; `ValueError` с именем файла при любой ошибке."""
    base = Path(directory)
    loaded: dict[str, BaseModel] = {}
    for name, model in ((LINES_FILE, _LinesFile), (LOOKALIKES_FILE, _LookalikesFile), (LABELS_FILE, _LabelsFile)):
        try:
            loaded[name] = model.model_validate(_read_yaml(base, name))
        except ValidationError as exc:
            raise ValueError(f"справочник {name}: {exc}") from exc
    lines = loaded[LINES_FILE]
    lookalikes = loaded[LOOKALIKES_FILE]
    labels = loaded[LABELS_FILE]
    assert isinstance(lines, _LinesFile) and isinstance(lookalikes, _LookalikesFile)  # noqa: S101 — сужение типов
    assert isinstance(labels, _LabelsFile)  # noqa: S101
    return OcrRules(editions=lines.editions, lookalikes=lookalikes.digits, labels=labels.labels)