medtech-hackathon/etl/extractors/readers.py

"""Чтение исходных файлов в «сетки» строк плюс метаданные документа.

Каждый читатель возвращает (список сеток, сырой текст). Несколько сеток бывает у
многолистовых книг и многотабличных PDF. Сырой текст нужен для аудита и для
детектора битого текстового слоя.
"""

from __future__ import annotations

import re
from pathlib import Path

from etl.extractors.common import Grid, clean


def year_from_name(name: str) -> int | None:
    """Год прайса из имени файла («Клиника 2 прайс 2025 год» → 2025)."""
    m = re.search(r"(20\d{2})", name)
    return int(m.group(1)) if m else None


def partner_from_name(name: str) -> str:
    """Название клиники из имени файла, очищенное от слов «прайс», года и т. п."""
    base = Path(name).stem
    base = re.sub(r"(?i)\b(прайс|price|прейскурант|год)\b", " ", base)
    base = re.sub(r"20\d{2}", " ", base)
    base = re.sub(r"[_\-]+", " ", base)
    return re.sub(r"\s+", " ", base).strip()


_KZ_CITIES = {
    "алмат": "Алматы", "астан": "Астана", "нур-султан": "Астана", "шымкент": "Шымкент",
    "караганд": "Караганда", "актобе": "Актобе", "павлодар": "Павлодар", "семей": "Семей",
    "усть-камен": "Усть-Каменогорск", "костанай": "Костанай", "атырау": "Атырау", "тараз": "Тараз",
    "уральск": "Уральск", "кокшетау": "Кокшетау", "петропавл": "Петропавловск",
    "кызылорд": "Кызылорда", "талдыкорган": "Талдыкорган", "актау": "Актау",
}


def city_from_text(text: str | None) -> str | None:
    """Город партнёра по тексту документа (БИН в казахстанских прайсах обычно не печатают)."""
    low = (text or "").lower()
    for stem, name in _KZ_CITIES.items():
        if stem in low:
            return name
    return None


def detect_format(path: str) -> str:
    ext = Path(path).suffix.lower()
    return {".xlsx": "xlsx", ".xls": "xls", ".docx": "docx", ".pdf": "pdf"}.get(
        ext, ext.lstrip(".")
    )


def docx_grids(path: str) -> tuple[list[Grid], str]:
    import docx

    document = docx.Document(path)
    grids = [
        [[clean(cell.text) for cell in row.cells] for row in table.rows]
        for table in document.tables
    ]
    text = "\n".join(p.text for p in document.paragraphs)
    return grids, text


def xlsx_grids(path: str) -> tuple[list[Grid], str]:
    from openpyxl import load_workbook

    workbook = load_workbook(path, read_only=True, data_only=True)
    grids = []
    for sheet in workbook.worksheets:
        grid = [[clean(c) for c in row] for row in sheet.iter_rows(values_only=True)]
        if any(any(row) for row in grid):
            grids.append(grid)
    return grids, ""


def xls_grids(path: str) -> tuple[list[Grid], str]:
    import xlrd

    book = xlrd.open_workbook(path)
    grids = []
    for sheet in book.sheets():
        grid = [
            [clean(sheet.cell_value(r, c)) for c in range(sheet.ncols)] for r in range(sheet.nrows)
        ]
        if any(any(row) for row in grid):
            grids.append(grid)
    return grids, ""


def _line_to_cells(words: list[dict], x_gap: float) -> list[str]:
    """Слова одной строки → ячейки: большой горизонтальный разрыв = граница колонки."""
    words.sort(key=lambda w: w["x0"])
    cells = [words[0]["text"]]
    right = words[0]["x1"]
    for w in words[1:]:
        if w["x0"] - right > x_gap:
            cells.append(w["text"])
        else:
            cells[-1] += " " + w["text"]
        right = w["x1"]
    return [clean(c) for c in cells]


def _page_to_grid(page, y_tol: float = 3.0, x_gap: float = 18.0) -> Grid:
    """Восстановить табличную сетку страницы PDF по координатам слов.

    У многих прайсов нет линий таблицы, поэтому pdfplumber.extract_tables ничего не
    находит. Здесь слова группируются в строки по вертикали (top), а внутри строки
    режутся на ячейки по горизонтальным разрывам — так колонки цен (резидент/
    нерезидент/…) не сливаются в одну.
    """
    words = page.extract_words(use_text_flow=False, keep_blank_chars=False)
    if not words:
        return []
    rows: Grid = []
    line: list[dict] = []
    current_band: int | None = None
    for w in sorted(words, key=lambda w: (round(w["top"] / y_tol), w["x0"])):
        band = round(w["top"] / y_tol)
        if current_band is None or band == current_band:
            line.append(w)
        else:
            rows.append(_line_to_cells(line, x_gap))
            line = [w]
        current_band = band
    if line:
        rows.append(_line_to_cells(line, x_gap))
    return rows


def pdf_grids(path: str) -> tuple[list[Grid], str]:
    """Все страницы PDF собираются в одну сетку — так шапка распознаётся один раз."""
    import pdfplumber

    grid: Grid = []
    texts: list[str] = []
    with pdfplumber.open(path) as pdf:
        for page in pdf.pages:
            texts.append(page.extract_text() or "")
            grid.extend(_page_to_grid(page))
    return [grid], "\n".join(texts)


def pdf_page_count(path: str) -> int:
    """Число страниц PDF — нужно для оценки выхода строк и стоимости Vision."""
    import pdfplumber

    with pdfplumber.open(path) as pdf:
        return len(pdf.pages)


def is_garbled(text: str, min_letters: int = 200) -> bool:
    """Битый текстовый слой: мало кириллицы или много латиницы-мусора.

    Русский прайс почти не содержит латинских букв (кроме редких брендов), поэтому
    их обилие — верный признак испорченной кодировки. На таком документе нативный
    парсинг бессмыслен, его добирает путь через Vision.
    """
    letters = [ch for ch in text if ch.isalpha()]
    if len(letters) < min_letters:
        return False
    cyrillic = sum(1 for ch in letters if "а" <= ch.lower() <= "я" or ch.lower() == "ё")
    latin = sum(1 for ch in letters if "a" <= ch.lower() <= "z")
    return cyrillic / len(letters) < 0.55 or latin / len(letters) > 0.18