doccano/backend/data_import/pipeline/labels.py

import abc
from typing import Any, Dict, Optional, Union

from pydantic import BaseModel, validator

from api.models import Project
from label_types.models import LabelType, CategoryType, SpanType
from labels.models import Category, Span, TextLabel as TL


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: Union[str, int]
    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
        )