doccano/backend/data_import/pipeline/labels.py

import abc
from typing import Any, Dict, Optional

from pydantic import BaseModel, validator

from label_types.models import CategoryType, LabelType, RelationType, SpanType
from labels.models import Category
from labels.models import Label as LabelModel
from labels.models import Relation, Span
from labels.models import TextLabel as TL
from projects.models import Project


class Label(BaseModel, abc.ABC):
    @abc.abstractmethod
    def has_name(self) -> bool:
        raise NotImplementedError()

    @property
    @abc.abstractmethod
    def name(self) -> str:
        raise NotImplementedError()

    @classmethod
    def parse(cls, obj: Any):
        raise NotImplementedError()

    @abc.abstractmethod
    def create_type(self, project: Project) -> Optional[LabelType]:
        raise NotImplementedError()

    @abc.abstractmethod
    def create(self, user, example, mapping, **kwargs) -> LabelModel:
        raise NotImplementedError

    def __hash__(self):
        return hash(tuple(self.dict()))


class CategoryLabel(Label):
    label: str

    @validator("label")
    def label_is_not_empty(cls, value: str):
        if value:
            return value
        else:
            raise ValueError("is not empty.")

    def has_name(self) -> bool:
        return True

    @property
    def name(self) -> str:
        return self.label

    @classmethod
    def parse(cls, obj: Any):
        if isinstance(obj, str):
            return cls(label=obj)
        elif isinstance(obj, int):
            return cls(label=str(obj))
        else:
            raise TypeError(f"{obj} is not str.")

    def create_type(self, project: Project) -> Optional[LabelType]:
        return CategoryType(text=self.label, project=project)

    def create(self, user, example, mapping: Dict[str, LabelType], **kwargs):
        return Category(user=user, example=example, label=mapping[self.label])


class SpanLabel(Label):
    id: int = -1
    label: str
    start_offset: int
    end_offset: int

    def has_name(self) -> bool:
        return True

    @property
    def name(self) -> str:
        return self.label

    @classmethod
    def parse(cls, obj: Any):
        if isinstance(obj, list) or isinstance(obj, tuple):
            columns = ["start_offset", "end_offset", "label"]
            obj = zip(columns, obj)
            return cls.parse_obj(obj)
        elif isinstance(obj, dict):
            return cls.parse_obj(obj)
        else:
            raise TypeError(f"{obj} is invalid type.")

    def create_type(self, project: Project) -> Optional[LabelType]:
        return SpanType(text=self.label, project=project)

    def create(self, user, example, mapping: Dict[str, LabelType], **kwargs):
        return Span(
            user=user,
            example=example,
            start_offset=self.start_offset,
            end_offset=self.end_offset,
            label=mapping[self.label],
        )


class TextLabel(Label):
    text: str

    def has_name(self) -> bool:
        return False

    @property
    def name(self) -> str:
        return self.text

    @classmethod
    def parse(cls, obj: Any):
        if isinstance(obj, str) and obj:
            return cls(text=obj)
        else:
            raise TypeError(f"{obj} is not str or empty.")

    def create_type(self, project: Project) -> Optional[LabelType]:
        return None

    def create(self, user, example, mapping, **kwargs):
        return TL(user=user, example=example, text=self.text)


class RelationLabel(Label):
    from_id: int
    to_id: int
    type: str

    def has_name(self) -> bool:
        return True

    @property
    def name(self) -> str:
        return self.type

    @classmethod
    def parse(cls, obj: Any):
        if isinstance(obj, dict):
            return cls.parse_obj(obj)
        else:
            raise TypeError(f"{obj} is not dict.")

    def create_type(self, project: Project) -> Optional[LabelType]:
        return RelationType(text=self.type, project=project)

    def create(self, user, example, mapping: Dict[str, LabelType], **kwargs):
        return Relation(
            user=user,
            example=example,
            type=mapping[self.type],
            from_id=kwargs["span_mapping"][self.from_id],
            to_id=kwargs["span_mapping"][self.to_id],
        )