Browse Source

Add replace method to label

pull/1310/head
Hironsan 3 years ago
parent
commit
d0fea9c7e5
1 changed files with 20 additions and 1 deletions
  1. 21
      app/api/views/upload/label.py

21
app/api/views/upload/label.py

@ -1,5 +1,5 @@
import abc
from typing import Any
from typing import Any, Dict
from pydantic import BaseModel
@ -19,6 +19,10 @@ class Label(BaseModel, abc.ABC):
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()))
@ -39,6 +43,10 @@ class CategoryLabel(Label):
return cls(label=obj)
raise TypeError(f'{obj} is not str.')
def replace(self, mapping: Dict[str, int]) -> 'Label':
label = mapping.get(self.label, self.label)
return CategoryLabel(label=label)
class OffsetLabel(Label):
label: str
@ -63,6 +71,14 @@ class OffsetLabel(Label):
else:
raise TypeError(f'{obj} is invalid type.')
def replace(self, mapping: Dict[str, int]) -> 'Label':
label = mapping.get(self.label, self.label)
return OffsetLabel(
label=label,
start_offset=self.start_offset,
end_offset=self.end_offset
)
class TextLabel(Label):
text: str
@ -80,3 +96,6 @@ class TextLabel(Label):
return cls(text=obj)
else:
raise TypeError(f'{obj} is not str.')
def replace(self, mapping: Dict[str, str]) -> 'Label':
return self
Loading…
Cancel
Save