doccano/backend/api/views/upload/label.py

import abc
from typing import Any, Dict, Union

from pydantic import BaseModel


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 replace(self, mapping: Dict[str, int]) -> 'Label':
        raise NotImplementedError

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


class CategoryLabel(Label):
    label: Union[str, 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, str):
            return cls(label=obj)
        raise TypeError(f'{obj} is not str.')

    def replace(self, mapping: Dict[str, int]) -> 'Label':
        label = mapping[self.label]
        return CategoryLabel(label=label)


class OffsetLabel(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 replace(self, mapping: Dict[str, int]) -> 'Label':
        label = mapping[self.label]
        return OffsetLabel(
            label=label,
            start_offset=self.start_offset,
            end_offset=self.end_offset
        )


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 replace(self, mapping: Dict[str, str]) -> 'Label':
        return self