lib/parser/ntr_fileparser/data_classes/parsed_document.py

"""
Модуль содержит класс для представления структуры документа.
"""

from dataclasses import dataclass, field
from typing import Any, Callable

from .parsed_formula import ParsedFormula
from .parsed_image import ParsedImage
from .parsed_meta import ParsedMeta
from .parsed_structure import ParsedStructure
from .parsed_table import ParsedTable
from .parsed_text_block import ParsedTextBlock


@dataclass
class ParsedDocument(ParsedStructure):
    """
    Документ, полученный в результате парсинга.
    """
    name: str = ""
    type: str = ""
    meta: ParsedMeta = field(default_factory=ParsedMeta)
    paragraphs: list[ParsedTextBlock] = field(default_factory=list)
    tables: list[ParsedTable] = field(default_factory=list)
    images: list[ParsedImage] = field(default_factory=list)
    formulas: list[ParsedFormula] = field(default_factory=list)
    
    def to_string(self) -> str:
        """
        Преобразует документ в строковое представление.
        
        Returns:
            str: Строковое представление документа.
        """
        result = [f"Документ: {self.name} (тип: {self.type})"]
        
        if self.paragraphs:
            result.append("\nПараграфы:")
            for p in self.paragraphs:
                result.append(p.to_string())
        
        if self.tables:
            result.append("\nТаблицы:")
            for t in self.tables:
                result.append(t.to_string())
        
        if self.images:
            result.append("\nИзображения:")
            for i in self.images:
                result.append(i.to_string())
        
        if self.formulas:
            result.append("\nФормулы:")
            for f in self.formulas:
                result.append(f.to_string())
        
        return "\n".join(result)
    
    def apply(self, func: Callable[[str], str]) -> None:
        """
        Применяет функцию ко всем строковым элементам документа.
        
        Args:
            func (Callable[[str], str]): Функция для применения к текстовым элементам.
        """
        self.name = func(self.name)
        self.type = func(self.type)
        
        # Применяем к параграфам
        for p in self.paragraphs:
            p.apply(func)
        
        # Применяем к таблицам
        for t in self.tables:
            t.apply(func)
        
        # Применяем к изображениям
        for i in self.images:
            i.apply(func)
        
        # Применяем к формулам
        for f in self.formulas:
            f.apply(func)
    
    def to_dict(self) -> dict[str, Any]:
        """
        Преобразует документ в словарь.
        
        Returns:
            dict[str, Any]: Словарное представление документа.
        """
        # Преобразуем тип в строку, если это объект FileType
        type_str = str(self.type) if not isinstance(self.type, str) else self.type
        
        result = {
            'name': self.name,
            'type': type_str,
            'meta': self.meta.to_dict(),
            'paragraphs': [p.to_dict() for p in self.paragraphs],
            'tables': [t.to_dict() for t in self.tables],
            'images': [i.to_dict() for i in self.images],
            'formulas': [f.to_dict() for f in self.formulas]
        }
        return result