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