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, SpanType
from labels.models import Category, 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(self, project: Project) -> Optional[LabelType]:
        raise NotImplementedError()

    @abc.abstractmethod
    def create_annotation(self, user, example, mapping):
        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(self, project: Project) -> Optional[LabelType]:
        return CategoryType(text=self.label, project=project)

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


class SpanLabel(Label):
    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(self, project: Project) -> Optional[LabelType]:
        return SpanType(text=self.label, project=project)

    def create_annotation(self, user, example, mapping: Dict[str, LabelType]):
        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(self, project: Project) -> Optional[LabelType]:
        return None

    def create_annotation(self, user, example, mapping):
        return TL(user=user, example=example, text=self.text)