lib/extractor/tests/custom_entity.py

from uuid import UUID

from ntr_text_fragmentation.models.linker_entity import (LinkerEntity,
                                                         register_entity)


@register_entity
class CustomEntity(LinkerEntity):
    """Пользовательский класс-наследник LinkerEntity для тестирования сериализации и десериализации."""
    
    def __init__(
        self,
        id: UUID,
        name: str,
        text: str,
        metadata: dict,
        custom_field1: str,
        custom_field2: int,
        in_search_text: str | None = None,
        source_id: UUID | None = None,
        target_id: UUID | None = None,
        number_in_relation: int | None = None,
        type: str = "CustomEntity"
    ):
        super().__init__(
            id=id, 
            name=name, 
            text=text, 
            metadata=metadata, 
            in_search_text=in_search_text,
            source_id=source_id,
            target_id=target_id,
            number_in_relation=number_in_relation,
            type=type
        )
        self.custom_field1 = custom_field1
        self.custom_field2 = custom_field2
    
    def deserialize(self, entity: LinkerEntity) -> 'CustomEntity':
        """Реализация метода десериализации для кастомного класса."""
        custom_field1 = entity.metadata.get('_custom_field1', '')
        custom_field2 = entity.metadata.get('_custom_field2', 0)
        
        # Создаем чистые метаданные без служебных полей
        clean_metadata = {k: v for k, v in entity.metadata.items() 
                        if not k.startswith('_')}
        
        return CustomEntity(
            id=entity.id,
            name=entity.name,
            text=entity.text,
            in_search_text=entity.in_search_text,
            metadata=clean_metadata,
            source_id=entity.source_id,
            target_id=entity.target_id,
            number_in_relation=entity.number_in_relation,
            custom_field1=custom_field1,
            custom_field2=custom_field2
        )
    
    @classmethod
    def deserialize(cls, entity: LinkerEntity) -> 'CustomEntity':
        """
        Классовый метод для десериализации.
        Необходим для работы с реестром классов.
        
        Args:
            entity: Сериализованная сущность
            
        Returns:
            Десериализованный экземпляр CustomEntity
        """
        custom_field1 = entity.metadata.get('_custom_field1', '')
        custom_field2 = entity.metadata.get('_custom_field2', 0)
        
        # Создаем чистые метаданные без служебных полей
        clean_metadata = {k: v for k, v in entity.metadata.items() 
                        if not k.startswith('_')}
        
        return CustomEntity(
            id=entity.id,
            name=entity.name,
            text=entity.text,
            in_search_text=entity.in_search_text,
            metadata=clean_metadata,
            source_id=entity.source_id,
            target_id=entity.target_id,
            number_in_relation=entity.number_in_relation,
            custom_field1=custom_field1,
            custom_field2=custom_field2
        )
    
    def __eq__(self, other):
        """Переопределяем метод сравнения для проверки равенства объектов."""
        if not isinstance(other, CustomEntity):
            return False
            
        # Используем базовое сравнение из LinkerEntity, которое уже учитывает поля связи
        base_equality = super().__eq__(other)
        
        # Дополнительно проверяем кастомные поля
        return (
            base_equality
            and self.custom_field1 == other.custom_field1
            and self.custom_field2 == other.custom_field2
        )