24 lines
590 B
Python
24 lines
590 B
Python
from enum import Enum
|
|
|
|
from image_prediction.formatter.formatters.key_formatter import KeyFormatter
|
|
|
|
|
|
class EnumFormatter(KeyFormatter):
|
|
def format_key(self, key):
|
|
return key.value if isinstance(key, Enum) else key
|
|
|
|
def transform(self, obj):
|
|
raise NotImplementedError
|
|
|
|
|
|
class ReverseEnumFormatter(KeyFormatter):
|
|
def __init__(self, enum):
|
|
self.enum = enum
|
|
self.reverse_enum = {e.value: e for e in enum}
|
|
|
|
def format_key(self, key):
|
|
return self.reverse_enum.get(key, key)
|
|
|
|
def transform(self, obj):
|
|
raise NotImplementedError
|