django enum custom field 만들기

enum 클래스

enumerate는 열거형이라고 불리며 고유한 상숫값에 연결된 기호 이름(멤버)의 집합입니다. python은 열거형을 지원하기 위해 enum 클래스가 존재합니다.

아래는 열거형을 만드는 예시 코드입니다.

class LanguageType(Enum):
    C = 'c'
    PYTHON = 'python'
    JAVA = 'JAVA'

django와의 호환성

하지만 python의 enum 클래스는 django orm과 호환성이 좋지 못합니다. 무슨 의미인지 아래 코드를 보면 알 수 있습니다.

class TestType(Enum):
    A_TYPE = 'a_type'
    B_TYPE = 'b_type'
    C_TYPE = 'c_type'

class TestModel(models.Model):
    id = models.AutoField(primary_key=True)
    type = models.CharField(max_length=10)
    class Meta:
        db_table = 't_test_model'

>>> t_model = TestModel.objects.get(id=1)
>>> type(t_model.type)
<class 'str'>

위와 같이 t_model의 type값이 string 타입인 것을 확인할 수 있습니다.

맞습니다. django orm은 Enum 타입을 그대로 데이터베이스로부터 불러올 수 없고, int나 string 같은 primitive 타입 값을 불러오게 됩니다.

그렇기 때문에 발생하는 오류가 존재합니다.

>>> t_model.type
'a_type'
>>> t_model.type == TestType.A_TYPE
False

python의 enum 비교는 같은 eum 타입만 비교할 수 있기 때문에 값이 같더라도 enum 타입이 아닌 이상 다르다고 처리해버립니다.

그렇기 때문에 아래와 같이 django model 값과 enum 값을 비교하기 위해서는 enum 타입의 value 속성을 이용해 비교해야합니다.

>>> t_model.type == TestType.A_TYPE**.value**  # 명시적으로 value 값을 비교해줘야 한다.
True

그렇다면 모든 코드들을 value 속성을 이용해서 비교해야 하는데, 사람이 코드를 개발하는 특성 상 실수가 발생하기 쉬운 형태입니다.

하지만 django custom field를 이용한다면 이를 해결할 수 있습니다.

django custom field

django 공식 문서를 보면 django model에 대한 field를 커스텀으로 생성할 수 있는 방법이 있습니다.

field 클래스는 데이터베이스로부터 데이터를 불러올 때, 혹은 데이터를 저장할 때 변환하는 로직을 가지고 있습니다. 이를 통해 Enum 클래스로 변환하는 로직을 구현할 수 있습니다.

TextEnumField

위의 TestType처럼 text 형식의 enum을 위한 custom field는 아래와 같습니다.

class TextEnumField(models.CharField):
    def __init__(self, *args, **kwargs):
        self.enum = kwargs.pop('enum')
        self.__enum_value_to_key_map = {item.value: item for item in self.enum}
        if not issubclass(self.enum, Enum):
            raise TypeError('enum 인자는 Enum 형식이여야 합니다.')
        super().__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super().deconstruct()
        kwargs['enum'] = self.enum
        return name, path, args, kwargs

    def to_python(self, value):
        if isinstance(value, self.enum):
            return value
        if value is None:
            return None
        val = self.__enum_value_to_key_map[value]
        if val is None:
            raise Exception('존재하지 않는 Enum 타입입니다.')
        return val

    def from_db_value(self, value, expression, connection):
        return self.to_python(value=value)

    def get_prep_value(self, value):
        if isinstance(value, self.enum):
            return value.value
        else:
            return value

    def validate(self, value, model_instance):
        super().validate(value, model_instance)
        if not isinstance(value, self.enum):
            raise ValidationError('잘못된 Enum 타입입니다.')

IntEnumField

int 형식의 enum을 위한 custom field는 아래와 같습니다.

class IntEnumField(models.IntegerField):
    def __init__(self, *args, **kwargs):
        self.enum = kwargs.pop('enum')
        self.__enum_value_to_key_map = {item.value: item for item in self.enum}
        if not issubclass(self.enum, IntEnum):
            raise TypeError('enum 인자는 IntEnum 형식이여야 합니다.')
        super().__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super().deconstruct()
        kwargs['enum'] = self.enum
        return name, path, args, kwargs

    def to_python(self, value):
        if isinstance(value, self.enum):
            return value
        if value is None:
            return None
        val = self.__enum_value_to_key_map[value]
        if val is None:
            raise Exception('존재하지 않는 IntEnum 타입입니다.')
        return val

    def from_db_value(self, value, expression, connection):
        value = self.to_python(value=value)

        if value is None:
            return None
        return value

    def get_prep_value(self, value):
        if isinstance(value, self.enum):
            return value.value
        else:
            return value

    def validate(self, value, model_instance):
        super().validate(value, model_instance)
        if not isinstance(value, self.enum):
            raise ValidationError('잘못된 Choice 타입입니다.')

※ django custom field에 대한 자세한 설명은 추후에 추가할 예정입니다.