[Django] [Field][EnumField] TextChoice 말고 Enum을 이용하여 필드를 구성해보자!

만들게 된 계기

---

장고에서는 기본적으로 모델 필드에서 사용할 수 있는 enum은 choice 라는 것이 있다.
하지만 필자가 다양한 프레임워크를 사용하면서 느낀 가장 불편한 점은 다음과 같다.

1. IDE에서 제공해주는 hint들을 사용할 수 없다. (Python에서도 Enum이 존재하는 데 이 때 '.'을 찍으면 해당 Enum에 존재하는 값들에 대한 힌트들을 얻을 수 있는데 TextChoice의 경우에는 튜플로 되어 있어서 이러한 힌트들을 사용할 수 없다.)

 

1. 명시적이지 않다. (밑의 글을 읽어보면 제가 만든 EnumField가 TextChoice Field 보다 더 보기 좋다는 것을 느낄 것이다.)

 

1. 사용하기 더 편하다.

구현 방법

from django.db.models import CharField


class EnumField(CharField):

    def __init__(self, enum, *args, **kwargs):
        self.enum = enum
        super().__init__(*args, **kwargs)

    def get_default(self):
        default = super().get_default()
        self.validate_enum(default)
        return default

    def to_python(self, value):
        return super().to_python(self.validate_enum(value))

    def get_prep_value(self, value):
        return super().get_prep_value(self.validate_enum(value))

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

    def validate_enum(self, value):
        for name, member in self.enum.__members__.items():
            if member == value:
                return value.value
            if member.value == value:
                return value
        raise AttributeError('Not Found Enum Member')

init

여기에서 enum 필드라는 것을 명시한 것을 볼 수가 있습니다.
즉 enum 필드에 대한 값을 정의하지 않으면 해당 필드는 TypeError를 내보내면서 반드시 정의하도록 요구합니다.

get_default

해당 메소드는 필드에서 제공하는 기본 필드로서 default 값을 세팅한 값을 가져올 때 쓰이는 메소드입니다.
오버라이딩 한 이유는 enum값에 대한 validation 하기 위해 했습니다.

to_python

해당 메소드 또한 오버라이딩한 메소드로서 enum에 대한 value 체크 및 변환 작업을 위해 재정의 했습니다.

deconstruct

위에서 init 메소드에서 제가 추가적으로 enum이라는 파라미터를 추가했는데 이를 사용하기 위해 deconstruct 메소드 kwargs에 해당 내용을 명시해야 합니다.
공식문서 deconstruct 설명

validate_enum

애플리케이션에서는 enum으로 사용하더라도 데이터베이스에서는 결과적으로 문자열이 저장되어야 하므로 이에 대한 변환작업이 필수적입니다.
이를 위해 추가적인 메소드를 만들어서 애플리케이션에서 enum 또는 해당 직접적인 값을 넣었을 때에는 반환을 해주고 그렇지 않았을 때에는 AttributeError를 일으키도록 작성했습니다.

TEST

    def test_enum필드_요구(self):
        with self.assertRaises(TypeError):
            EnumField()

    def test_default_값_체크(self):
        field = EnumField(enum=TestEnum, default="hello")
        with self.assertRaises(AttributeError):
            field.get_default()

    def test_default_값_정상확인(self):
        field = EnumField(enum=TestEnum, default=TestEnum.UNIT)
        value = field.get_default()
        assert_that(value).is_equal_to(TestEnum.UNIT)

    def test_to_python_method(self):
        field = EnumField(enum=TestEnum, default=TestEnum.UNIT)
        value = field.to_python(value=TestEnum.TEST)
        assert_that(value).is_equal_to(TestEnum.TEST.value)

    def test_validate_enum(self):
        field = EnumField(enum=TestEnum)
        with self.assertRaises(AttributeError):
            field.validate_enum(AnotherTestEnum.YOU)

class TestEnum(Enum):
    TEST = 'Test'
    UNIT = 'Unit'

class AnotherTestEnum(Enum):
    YOU = 'YOU'
    AND = 'AND'

Model에서 사용법

class MyModel(Model):
    test_type = EnumField(enum=TestType, default=TestType.TDD, max_length=64)

필자는 사용하면서 Django에서 제공해주는 TextChoice 보다 훨씬 더 편하게 사용하였기에 이렇게 블로그로 해당 내용을 공유하고자 합니다.