components/parser/abbreviations/abbreviation_extractor.py

import re

from natasha import Doc, MorphVocab, NewsEmbedding, NewsMorphTagger, Segmenter


from .constants import (
    ABBREVIATION_RE,
    CLOSE_BRACKET_RE,
    FIRST_CHARS_SET,
    NEXT_MARKER_RE,
    NON_SENTENCE_ENDINGS,
    SECOND_CHARS_SET,
    UPPERCASE_LETTER_RE,
)
from .structures import Abbreviation


class AbbreviationExtractor:
    def __init__(self):
        """
        Инициализация экстрактора сокращений.

        Создает необходимые компоненты для лемматизации и компилирует регулярные выражения.
        """
        # Инициализация компонентов Natasha для лемматизации
        self.segmenter = Segmenter()
        self.morph_tagger = NewsMorphTagger(NewsEmbedding())
        self.morph_vocab = MorphVocab()

        # Компиляция регулярных выражений
        self.next_re = re.compile(NEXT_MARKER_RE, re.IGNORECASE)
        self.abbreviation_re = re.compile(ABBREVIATION_RE)
        self.uppercase_letter_re = re.compile(UPPERCASE_LETTER_RE)
        self.close_bracket_re = re.compile(CLOSE_BRACKET_RE)

        self.delimiters = [
            f'{char1} {char2} '.format(char1, char2)
            for char1 in FIRST_CHARS_SET
            for char2 in SECOND_CHARS_SET
        ]

    def extract_abbreviations_from_text(
        self,
        text: str,
    ) -> list[Abbreviation]:
        """
        Извлечение всех сокращений из текста.

        Args:
            text: Текст для обработки

        Returns:
            list[Abbreviation]: Список найденных сокращений
        """
        sentences = self._extract_sentences_with_abbreviations(text)

        abbreviations = [self._process_one_sentence(sentence) for sentence in sentences]
        abbreviations = sum(abbreviations, [])  # делаем список одномерным
        abbreviations = [abbreviation.process() for abbreviation in abbreviations]

        return abbreviations

    def _process_one_sentence(self, sentence: str) -> list[Abbreviation]:
        """
        Обработка одного предложения для извлечения сокращений.

        Args:
            sentence: Текст для обработки

        Returns:
            list[Abbreviation]: Список найденных сокращений
        """
        search_iter = self.next_re.finditer(sentence)
        prev_index = 0

        abbreviations = []

        for match in search_iter:
            abbreviation, prev_index = self._process_match(sentence, match, prev_index)
            if abbreviation is not None:
                abbreviations.append(abbreviation)

        return abbreviations

    def _process_match(
        self,
        sentence: str,
        match: re.Match,
        prev_index: int,
    ) -> tuple[Abbreviation | None, int]:
        """
        Обработка одного совпадения с конструкцией "далее - {short_form}" для извлечения сокращений.

        Args:
            sentence: Текст для обработки
            match: Совпадение для обработки
            prev_index: Предыдущий индекс

        Returns:
            tuple[Abbreviation | None, int]: Найденное сокращение (None, если нет сокращения) и следующий индекс
        """
        start, end = match.start(), match.end()
        text = sentence[start:]

        index_close_parenthesis = self._get_close_parenthesis_index(text)
        index_point = self._get_point_index(text, start)

        prev_index += index_point
        short_word = text[end : start + index_close_parenthesis].strip()

        if len(short_word.split()) < 2:
            abbreviation = self._process_match_for_word(
                short_word, text, start, end, prev_index
            )

        else:
            abbreviation = self._process_match_for_phrase(
                short_word, text, start, end, prev_index
            )

        prev_index = start + index_close_parenthesis + 1

        return abbreviation, prev_index

    def _get_close_parenthesis_index(self, text: str) -> int:
        """
        Получение индекса закрывающей скобки в тексте.

        Args:
            text: Текст для обработки

        Returns:
            int: Индекс закрывающей скобки или 0, если не найдено
        """
        result = self.close_bracket_re.search(text)
        if result is None:
            return 0
        return result.start()

    def _get_point_index(self, text: str, start_index: int) -> int:
        """
        Получение индекса точки в тексте.

        Args:
            text: Текст для обработки
            start_index: Индекс начала поиска

        Returns:
            int: Индекс точки или 0, если не найдено
        """
        result = text.rfind('.', 0, start_index - 1)
        if result == -1:
            return 0
        return result

    def _process_match_for_word(
        self,
        short_word: str,
        text: str,
        start_next_re_index: int,
        end_next_re_index: int,
        prev_index: int,
    ) -> Abbreviation | None:
        """
        Обработка сокращения, состоящего из одного слова.

        Args:
            short_word: Сокращение
            text: Текст для обработки
            start_next_re_index: Индекс начала следующего совпадения
            end_next_re_index: Индекс конца следующего совпадения
            prev_index: Предыдущий индекс

        Returns:
            Abbreviation | None: Найденное сокращение или None, если нет сокращения
        """
        if self.abbreviation_re.findall(text) or (short_word == 'ПДн'):
            return None

        lemm_text = self._lemmatize_text(text[prev_index:start_next_re_index])
        lemm_short_word = self._lemmatize_text(short_word)

        search_word = re.search(lemm_short_word, lemm_text)

        if not search_word:
            start_text_index = self._get_start_text_index(
                text,
                start_next_re_index,
                prev_index,
            )

            if start_text_index is None:
                return None

            full_text = text[prev_index + start_text_index : end_next_re_index]

        else:
            index_word = search_word.span()[1]
            space_index = text[prev_index:start_next_re_index].rfind(' ', 0, index_word)
            if space_index == -1:
                space_index = 0
            text = text[prev_index + space_index : start_next_re_index]

        full_text = text.replace(')', '').replace('(', '').replace('', '- ')

        return Abbreviation(
            short_form=short_word,
            full_form=full_text,
        )

    def _process_match_for_phrase(
        self,
        short_word: str,
        text: str,
        start_next_re_index: int,
        end_next_re_index: int,
        prev_index: int,
    ) -> list[Abbreviation] | None:
        """
        Обработка сокращения, состоящего из нескольких слов.
        В действительности производится обработка первого слова сокращения, а затем вместо него подставляется полное сокращение.

        Args:
            short_word: Сокращение
            text: Текст для обработки
            start_next_re_index: Индекс начала следующего совпадения
            end_next_re_index: Индекс конца следующего совпадения
            prev_index: Предыдущий индекс

        Returns:
            list[Abbreviation] | None: Найденные сокращения или None, если нет сокращений
        """
        first_short_word = short_word.split()[0]
        result = self._process_match_for_word(
            first_short_word, text, start_next_re_index, end_next_re_index, prev_index
        )
        if result is None:
            return None
        return Abbreviation(
            short_form=short_word,
            full_form=result.full_form,
        )

    def _get_start_text_index(
        self,
        text: str,
        start_next_re_index: int,
        prev_index: int,
    ) -> int | None:
        """
        Получение индекса начала текста для поиска сокращения с учётом разделителей типа
        "; - "
        ": - "
        ";  "
        ": ‒ " и т.п.

        Args:
            text: Текст для обработки
            start_next_re_index: Индекс начала следующего совпадения
            prev_index: Предыдущий индекс

        Returns:
            int | None: Индекс начала текста или None, если не найдено
        """
        if prev_index == 0:
            return 0

        for delimiter in self.delimiters:
            result = re.search(delimiter, text[prev_index:start_next_re_index])
            if result is not None:
                return result.span()[1]

        return None

    def _lemmatize_text(self, text: str) -> str:
        """
        Лемматизация текста.

        Args:
            text: Текст для лемматизации

        Returns:
            str: Лемматизированный текст
        """
        doc = Doc(text)
        doc.segment(self.segmenter)
        doc.tag_morph(self.morph_tagger)

        for token in doc.tokens:
            token.lemmatize(self.morph_vocab)

        return ' '.join([token.lemma for token in doc.tokens])

    def _extract_sentences_with_abbreviations(self, text: str) -> list[str]:
        """
        Разбивает текст на предложения с учетом специальных сокращений.

        Точка после сокращений из NON_SENTENCE_ENDINGS не считается концом предложения.

        Args:
            text: Текст для разбиения

        Returns:
            list[str]: Список предложений
        """
        text = text.replace('\n', ' ')
        sentence_endings = re.finditer(r'\.\s+[А-Я]', text)

        sentences = []
        start = 0

        for match in sentence_endings:
            end = match.start() + 1

            # Проверяем, не заканчивается ли предложение на специальное сокращение
            preceding_text = text[start:end]
            words = preceding_text.split()

            if words and any(
                words[-1].rstrip('.').startswith(abbr) for abbr in NON_SENTENCE_ENDINGS
            ):
                continue

            sentence = text[start:end].strip()
            sentences.append(sentence)
            start = end + 1

        # Добавляем последнее предложение
        if start < len(text):
            sentences.append(text[start:].strip())

        return [
            sentence
            for sentence in sentences
            if self.next_re.search(sentence) is not None
        ]