drhider/drhider/replacer.py

"""
Проход 2: замена сущностей в документах.

Два метода:
1. apply_replacements — замена в Markdown/plain-тексте (основной)
2. replace_in_docx — замена в DOCX с сохранением форматирования
   (для опционального DOCX-выхода, будет добавлен позже)
"""

import io
import re
from typing import Dict, List


def _md_tolerant_pattern(original: str) -> str:
    r"""Паттерн, допускающий ** между частями строки (Markdown bold).

    DOCX → Markdown: жирность добавляет ** вокруг слов.
    Если ** попадает внутрь сущности (например, **ООО **"**НУБЕС**"),
    точное совпадение не срабатывает.

    Этот паттерн вставляет (?:\*{1,2})? между каждым словом/символом,
    позволяя найти сущность даже с ** внутри.

    Args:
        original: Исходная строка сущности (например, 'ООО "НУБЕС"')

    Returns:
        Regex-паттерн, находящий ту же строку с ** между частями
    """
    MD = r'(?:\*{1,2})?'
    parts = re.findall(r'[А-ЯЁа-яёA-Za-z0-9]+|[^А-ЯЁа-яёA-Za-z0-9]', original)
    if len(parts) <= 1:
        return None  # Одно слово — ** не может разорвать изнутри
    return MD + MD.join(re.escape(p) for p in parts) + MD


def _build_combined_re(sorted_keys: List[str]):
    """Построить единый regex из всех ключей (убывание длины, границы слова).

    Возвращает None, если нет ни одного применимого ключа.
    """
    alternatives = []
    for original in sorted_keys:
        if not original:
            continue
        esc = re.escape(original)
        # Границы слова — только если сущность обрамлена буквами/цифрами
        if original[0].isalnum() and original[-1].isalnum():
            alternatives.append(r'(?<!\w)' + esc + r'(?!\w)')
        else:
            alternatives.append(esc)

    if not alternatives:
        return None
    return re.compile('|'.join(alternatives))


def apply_replacements(text: str, mapping: Dict[str, str], sorted_keys: List[str],
                       compiled_re=None) -> str:
    """Применить все замены из словаря mapping к строке.

    Ключи применяются в порядке убывания длины (sorted_keys).
    Это гарантирует, что более длинные совпадения заменяются раньше
    коротких (например, «ИНН 123456789012» до «ИНН 1234567890»).

    Для сущностей, начинающихся и заканчивающихся на букву/цифру,
    используются границы слова (\\b), чтобы избежать частичных замен.

    Args:
        text: Исходный текст
        mapping: Словарь {оригинал: замена}
        sorted_keys: Ключи mapping, отсортированные по длине (убывание)
        compiled_re: Опционально, заранее скомпилированный regex из _build_combined_re
                     (кэш между файлами). Если None — компилируется здесь.

    Returns:
        Текст с заменами
    """
    result = text

    # ── Однопроходная замена: единый regex из всех ключей ──
    # Вместо N отдельных re.sub по всему тексту (O(N×text)) собираем
    # один паттерн из всех ключей и делаем один re.sub с callback (O(text)).
    if not sorted_keys:
        return result

    combined = compiled_re if compiled_re is not None else _build_combined_re(sorted_keys)
    if combined is None:
        return result

    matched_keys = set()

    def _replace(match):
        m = match.group(0)
        matched_keys.add(m)
        return mapping.get(m, m)

    result = combined.sub(_replace, result)

    # ── Фоллбэк: Markdown-жирность могла разорвать строку ──
    # Применяется только к ключам, которые основная замена НЕ нашла
    # (например, «ООО **"**НУБЕС**"» — ** внутри сущности).
    for original in sorted_keys:
        if not original or original in matched_keys:
            continue
        md_pattern = _md_tolerant_pattern(original)
        if md_pattern is not None:
            result = re.sub(md_pattern, mapping[original], result)

    return result


def replace_in_docx(doc, mapping: Dict[str, str], sorted_keys: List[str]) -> bytes:
    """Заменить сущности в DOCX-документе с сохранением форматирования.

    Алгоритм для каждого параграфа:
    1. Склеиваем текст всех runs в одну строку
    2. Применяем замены
    3. Пишем результат в первый run, очищаем остальные

    Это нужно потому что python-docx разбивает текст на runs
    (например, mid-docx форматирование), и сущность может быть
    разорвана между несколькими runs.

    Args:
        doc: Объект python-docx Document
        mapping: Словарь замен
        sorted_keys: Ключи mapping по убыванию длины

    Returns:
        Бинарное содержимое изменённого DOCX-файла
    """
    # ── Обработка параграфов ──
    for para in doc.paragraphs:
        if not para.runs:
            continue
        # Склеиваем все runs в одну строку
        full_text = "".join(run.text for run in para.runs)
        replaced = apply_replacements(full_text, mapping, sorted_keys)

        # Если были замены — пишем в первый run, очищаем остальные
        if replaced != full_text:
            para.runs[0].text = replaced
            for run in para.runs[1:]:
                run.text = ""

    # ── Обработка таблиц ──
    for table in doc.tables:
        for row in table.rows:
            for cell in row.cells:
                for para in cell.paragraphs:
                    if not para.runs:
                        continue
                    full_text = "".join(run.text for run in para.runs)
                    replaced = apply_replacements(full_text, mapping, sorted_keys)

                    if replaced != full_text:
                        para.runs[0].text = replaced
                        for run in para.runs[1:]:
                            run.text = ""

    # Сохраняем в буфер
    buf = io.BytesIO()
    doc.save(buf)
    return buf.getvalue()