[python] Enum + Enum 커스터마이징

2024년 3월 7일 경 Enum에 대해서 알게 되어 velog에 포스팅 한 바 있다.

Enum 포스팅 : https://velog.io/@heyggun/Python-Enum

얼추 이해했다고 생각했는데, 자꾸 실무에서 핸들링 할 때 헷갈리는 내 자신 발견.
다시 한번 Enum 뽀개기 위해서 공부하면서 기록한다.

1. Basic Concept.

알아야할 개념

열거형 클래스 (Enumeration Class)

열거형(enum)은 관련된 상수 값들을 묶어서 이름 붙여 관리 하는 것 이다.
그냥 값만 나열하는 것이 아니라, 의미 있는 이름으로 관리해서 코드의 가독성과 안정성을 높인다.
python에서는 enum 모듈의 Enum 클래스를 이용해서 열거형 클래스를 만들 수 있다.

예를 들어 사람의 성별을 문자로 표현한다고 할 때

GENDER = "male"

이라고 할당할 때, 오타로 mael 이렇게 들어갈 수도 있다. 이를 막기 위해서 미리 정해진 값으로만 쓸 수 있게 하는 것이 열거형이다.

열거형으로 하면

from enum import Enum

class Gender(Enum):
	MALE = "male"
    FEMALE = "female"
    
gender = Gender.MALE

이렇게 Gender 클래스 안에 정의된 값만 쓸 수 있고, 잘못된 값을 코드 작성 시 부터 막을 수 있다.

[예시 1]

from enum import Enum

class Color(Enum):
    RED = "RED"
    BLUE = "BLUE"
    GREEN = "GREEN"
    
print(Color.RED.name) # 'RED'
print(Color.RED.value) # 'RED'
print(Color.RED) # Color.RED (객체)
print(type(Color.RED)) # <enum 'Color'>
print(str(Color.RED)) # 'Color.RED'
print(type(str(Color.RED))) # <class 'str'>
print(Color.RED == "RED") # False

[예시 2]

from enum import Enum

class Color(Enum):
    RED = 1
    BLUE = 2
    GREEN = 3
    
print(Color.RED.name) # 'RED'
print(Color.RED.value) # 1
print(Color.RED) # Color.RED (객체)
print(type(Color.RED)) # <enum 'Color'>
print(str(Color.RED)) # 'Color.RED'
print(type(str(Color.RED))) # <class 'str'>
print(Color.RED == "RED") # False

위와 같이 python에서 Enum 멤버는 기본적으로 str, int 같은 형을 따르지 않아서 그냥 Color.RED 라고 하면 Color.RED라는 객체로 동작한다.

[예시 1]

from enum import Enum

class Color(str, Enum):
    RED = "RED"
    BLUE = "BLUE"
    GREEN = "GREEN"
    
print(Color.RED.name) # 'RED'
print(Color.RED.value) # 'RED'
print(Color.RED) # Color.RED (객체)
print(type(Color.RED)) # <enum 'Color'>
print(str(Color.RED)) # 'Color.RED'
print(type(str(Color.RED))) # <class 'str'>
print(Color.RED == "RED") # True

[예시 2]

from enum import Enum

class Color(str, Enum):
    RED = 1
    BLUE = 2
    GREEN = 3
    
print(Color.RED.name) # 'RED'
print(Color.RED.value) # 1
print(Color.RED) # Color.RED (객체)
print(type(Color.RED)) # <enum 'Color'>
print(str(Color.RED)) # 'Color.RED'
print(type(str(Color.RED))) # <class 'str'>
print(Color.RED == "RED") # False
print(Color.RED == "1") # True

그런데 str, Enum을 같이 상속하면 Color.RED가 str도 상속 받은 객체 라서 "RED"라는 문자열과 동등 비교 할 때 True가 된다.

여기까지는 완전히 이해 완.

---

추가로 열거형 클래스로 아래와 같이 정의해서 핸들링 해서 사용하려고 했다.

class PersonalInfoCategory(str, Enum):
    UNIQUE_ID = ("UNIQUE_ID", "고유식별정보")
    FINANCIAL = ("FINANCIAL", "금융정보")
    GENERAL = ("GENERAL", "기타개인정보")
    ETC = ("ETC", "기타민감정보")

but,

TypeError: decoding str is not supported

TypeError가 났는데, str, Enum 조합을 사용할 때 기본적으로 멤버의 값은 str 이여야 한다. 나의 코드는 tuple 형태였는데, 그래서 str.__new__ 할 때 튜플을 str로 변환할 수 없어서 TypeError가 뜨는 것이다.

자 그렇다고 하면 내 입맛대로 핸들링하려면 어떤 방법을 써야 할까?

2. Enum을 커스텀해서 속성/값을 다루는 방법

1. Enum만 상속 + value를 복합형으로 두고 @property로 접근
2. str, Enum (또는 int, Enum) 상속 + __new__로 속성 추가
3. Enum만 상속 + __init__으로 속성 추가

등의 방법이 있다.

위 세 가지 방법을 통해서 아래와 같이 사용해본다.
이제 여기 PersonalInfoCategory라는 클래스가 있다.

[case 1] Enum만 상속 + value를 복합형으로 두고 @property로 접근 할 때

from enum import Enum

class PersonalInfoCategory(Enum):
    UNIQUE_ID = ("UNIQUE_ID", "고유식별정보")
    FINANCIAL = ("FINANCIAL", "금융정보")
    GENERAL = ("GENERAL", "기타개인정보")
    ETC = ("ETC", "기타민감정보")

    @property
    def code(self):
        return self.value[0]

    @property
    def display_name(self):
        return self.value[1]

value를 tuple(튜플)로 가져가면서 @property로 튜플 내부 값을 깔끔하게 접근할 수 있다.
장점으로는 직관적인 코드, value 타입 자유(튜플, 딕셔너리 등 복합형 가능), 직렬화 로직 커스텀이 가능하다는 것이다.
단점으로는 직렬화 할 때 value가 튜플이라 JSON 변환 시 불편한 것. Enum의 멤버를 str과 직접 비교할 수 없는 것이다. (=="UNIQUE_ID" 안됨) -> 그대신 위에서 property로 정의 했으니 code로 꺼내와서 비교할 수 있음.

이해를 돕기 위해 예시 코드를 좀 보자면

# 객체 직렬화

data = {"category" : PersonalInfoCategory.UNIQUE_ID}

print(data['category']) # PersonalInfoCategory.UNIQUE_ID
print(data['category'].code) # UNIQUE_ID
print(data['category'].display_name) # 고유식별정보

json으로 직렬화를 시도 할 때,

# JSON 직렬화
import json

data = {"category" : PersonalInfoCategory.UNIQUE_ID}

json_data = json.dumps(data)

TypeError: Object of type PersonalInfoCategory is not JSON serializable 가 발생한다.

이 방법은 내부 로직에서만 쓰고 직렬화가 필요 없는 경우나 복합형 데이터가 필요할 때 사용 가능하다.

뭐 직렬화는 아래와 같은 방법으로 가능

#커스텀 직렬화

data = {"category" : PersonalInfoCategory.UNIQUE_ID}

json_data = json.dumps({"category" : data["category"].code,
                        "display_name" : data["category"].display_name},
                       ensure_ascii=False)

json_data

매번 code를 꺼내줘야 하는 불편함이 있다.

if PersonalInfoCategory.UNIQUE_ID == "UNIQUE_ID":
    print("같은 값")
    
else:
    print("다른 값")

에서는 "다른 값"이 출력된다.

if PersonalInfoCategory.UNIQUE_ID.code == "UNIQUE_ID":
    print("같은 값")
    
else:
    print("다른 값")

에서는 "같은 값" 이 출력된다.

그런데 사실 이런 방법 외에도 namedtuple이나 dataclass를 value로 넣는 방법도 있다.

namedtuple을 value로 넣는 방법

from enum import Enum
from collections import namedtuple

Category = namedtuple("Category", ["code", "display_name"])

class PersonalInfoCategory(Enum):
    UNIQUE_ID = Category("UNIQUE_ID", "고유식별정보")
    FINANCEIAL = Category("FINANCIAL", "금융정보")
    GENERAL = Category("GENERAL", "기타개인정보")
    ETC = Category("ETC", "기타민감정보")
    
    @property
    def code(self):
        return self.value.code
    
    @property
    def display_name(self):
        return self.value.display_name
 
 
category = PersonalInfoCategory.UNIQUE_ID
print(category.code) # UNIQUE_ID
print(category.display_name) # 고유식별정보

# Enum 멤버 순회하기

for category in PersonalInfoCategory:
    print(category.code, category.display_name)

# UNIQUE_ID 고유식별정보
# FINANCIAL 금융정보
# GENERAL 기타개인정보
# ETC 기타민감정보

dataclass를 value로 넣는 방법

from enum import Enum
from dataclasses import dataclass

@dataclass(frozen=True)
class Category:
    code:str
    display_name:str
    

class PersonalInfoCategory(Enum):
    UNIQUE_ID = Category("UNIQUE_ID", "고유식별정보")
    FINANCIAL = Category("FINANCIAL", "금융정보")
    GENERAL = Category("GENERAL", "기타개인정보")
    ETC = Category("ETC", "기타민감정보")
    
    @property
    def code(self):
        return self.value.code
    
    @property
    def display_name(self):
        return self.value.display_name


print(PersonalInfoCategory.UNIQUE_ID.code) # UNIQUE_ID
print(PersonalInfoCategory.UNIQUE_ID.display_name) # 고유식별정보

for category in PersonalInfoCategory:
    print(category.code, category.display_name)
    
# UNIQUE_ID 고유식별정보
# FINANCIAL 금융정보
# GENERAL 기타개인정보
# ETC 기타민감정보

# 데이터 직렬화
from dataclasses import asdict
import json

category = PersonalInfoCategory.UNIQUE_ID
category_json = json.dumps(asdict(category.value), ensure_ascii=False)
print(category_json)

# {"code": "UNIQUE_ID", "display_name": "고유식별정보"}

data = [asdict(category.value) for category in PersonalInfoCategory]
data_json = json.dumps(data, ensure_ascii=False)
print(data_json)

# [{"code": "UNIQUE_ID", "display_name": "고유식별정보"}, {"code": "FINANCIAL", "display_name": "금융정보"}, {"code": "GENERAL", "display_name": "기타개인정보"}, {"code": "ETC", "display_name": "기타민감정보"}]

애초에 사실 위 클래스를 정의 할때 to_dict_list를 classmethod로 정의하는 방법도 있다.

from dataclasses import asdict
import json


@dataclass
class Category:
    code:str
    display_name:str
    
class PersonalInfoCategory(Enum):
    UNIQUE_ID = Category("UNIQUE_ID", "고유식별정보")
    FINANCIAL = Category("FINANCIAL", "금융정보")
    GENERAL = Category("GENERAL", "기타개인정보")
    ETC = Category("ETC", "기타민감정보")
    
    @property
    def code(self):
        return self.value.code
    
    @property
    def display_name(self):
        return self.value.display_name
    
    @classmethod
    def to_dict_list(cls):
        return [asdict(member.value)for member in cls]
    
    
data = PersonalInfoCategory.to_dict_list()
print(data)

# [{'code': 'UNIQUE_ID', 'display_name': '고유식별정보'}, {'code': 'FINANCIAL', 'display_name': '금융정보'}, {'code': 'GENERAL', 'display_name': '기타개인정보'}, {'code': 'ETC', 'display_name': '기타민감정보'}]


json_data = json.dumps(data, ensure_ascii=False)
print(json_data)

# [{"code": "UNIQUE_ID", "display_name": "고유식별정보"}, {"code": "FINANCIAL", "display_name": "금융정보"}, {"code": "GENERAL", "display_name": "기타개인정보"}, {"code": "ETC", "display_name": "기타민감정보"}]

[case 2] str, Enum + __new__ 방식

아래의 str,Enum 방식은 이러한 방법없이 바로 가능해서 직렬화 중심의 실무에서 많이 쓰인다. 아래를 보자.

class PersonalInfoCategory(str, Enum):
    UNIQUE_ID = ("UNIQUE_ID", "고유식별정보")
    FINANCIAL = ("FINANCIAL", "금융정보")
    GENERAL = ("GENERAL", "기타개인정보")
    ETC = ("ETC", "기타민감정보")

    def __new__(cls, value, display_name):
        obj = str.__new__(cls, value)
        obj._value_ = value
        obj.display_name = display_name
        return obj

위와 같이 클래스를 지정한 경우.

TypeError: _value_ not set in __new__, unable to create it 역시나 typeError가 발생한다.
str, Enum을 같이 상속할 때 __new__ 메서드를 서서 value랑 부가 속성을 설정할 때에도 __new__는 반드시 str 타입 인자만 받는다.

하고 싶다면

from enum import Enum

class PersonalInfoCategory(str, Enum):
    UNIQUE_ID = "UNIQUE_ID"
    FINANCIAL = "FINANCIAL"
    GENERAL = "GENERAL"
    ETC = "ETC"

    def __new__(cls, value):
        obj = str.__new__(cls, value)
        obj._value_ = value
        return obj

로만 쓸 수 있는 것이다.
내가 희망하는건 이게 아닌데!
하지만 결론은 str 상속 + Enum 에서는 value로 튜플을 못쓴다. (그렇게 됐다.)

그 원인은

Enum은 생성할 때 __new__(cls, value) 를 호출하면서 str.__new__(cls, value)를 호출하는데 여기서 value가 튜플이면 str로 변환할 수 없으니까 터진다.

그래서 str, Enum 상속시 value는 무조건 단일 string 값.
추가 정보는 property로 구현한다. (dict-like string은 추천 안함)

[case 3] Enum + __init__ 속성 추가

from enum import Enum

class PersonalInfoCategory(Enum):
    UNIQUE_ID = "UNIQUE_ID"
    FINANCIAL = "FINANCIAL"
    GENERAL = "GENERAL"
    ETC = "ETC"
    
    def __init__(self, value):
        self.display_name = {
            "UNIQUE_ID": "고유식별정보",
            "FINANCIAL": "금융정보",
            "GENERAL": "기타개인정보",
            "ETC": "기타민감정보"
        }[value]
    

value가 str이 되면서 __init__에서 value 값으로 추가 속성을 매핑하는 방법이다. 구현이 심플하고 직렬화시 value가 str이다. 그러나 value 값 하나만 가능하고, 속성값을 추가하려면 딕셔너리에 계속 추가해야 한다. Enum의 개수가 늘어나면 관리가 복잡해진다.

---

아무튼 결론은 살펴보니까 dataclass랑 Enum을 조합해서 사용하고, 각자 property 를 사용하고, 직렬화는 classmethod로 만들어 놓는게 젤 가독성도 좋고 핸들링이 좋아서 이렇게 가기로 했다.