파이썬 - Enum

파이썬 3.4 부터 Enum 이 추가되었다. 기존에는 Enum 을 직접 구현하여 썻었는데, 추가됨에 따라 편리해졌다.

 

Enum 을 사용하면, 범위에 대한 경계선을 만들 수 있다.

이를 통해, 코드의 가독성도 높일 수 있다.

또한 코드 완성(힌트) 를 통해 코드 작성 시 생길 수 있는 에러도 줄일 수 있다.

기본

from enum import Enum

class Fruit(Enum):
    APPLE = 1
    BANANA = 2
    ORANGE = 3

클래스 형태로 위와 같이 Enum 을 사용할 수 있다.

>>> Fruits.APPLE
<Fruits.APPLE: 1>
>>> Fruits.APPLE.name
'APPLE'
>>> Fruits.APPLE.value
1

클래스에서 선언한 열거형 멤버는 Fruits.APPLE 로 접근하여 사용할 수 있다.

열거형 멤버는 name 과 value 를 가지고 있다.

Fruits.APPLE 은 name 으로 'APPLE' 를 가지고 있고, value 로 1 을 가지고 있다.

name 과 value는 Fruits.APPLE.name 과 Fruits.APPLE.value 로 사용할 수 있다.

 

추가로 열거형에 간단히 값을 할당하는 방법으로는 아래 두 방법이 있다.

from enum import Enum, auto

# auto 를 통해 자동 할당.
class Fruits(Enum):
    APPLE = auto()
    BANANA = auto()
    ORANGE = auto()

# range 를 통한 할당
class Fruits(Enum):
	APPLE, BANANA, ORANGE = range(1,4)

값을 자동으로 할당하고 싶다면 auto 를 사용하면 된다. 1 부터 할당된다.

range 를 통해서도 값을 할당할 수도 있다.

 

for 문

>>> for item in Fruits:
...     print(item)
Fruits.APPLE
Fruits.BANANA
Fruits.ORANGE

for 문을 사용하여 반복문에 사용할 수 있다. 반복시 열거형 멤버자체가 반환된다.

열거형 멤버 확인

>>> print(Fruits.__members__)
{'APPLE': <Fruits.APPLE: 1>, 'BANANA': <Fruits.BANANA: 2>, 'ORANGE': <Fruits.ORANGE: 3>}

__member__ 를 통해 각 열거형 멤버가 어떤 것이 있는지 파악할 수 있다.

간편 사용방법

class Fruits(Enum):
    APPLE = auto()
    BANANA = auto()
    ORANGE = auto()

    def __eq__(self, value):
        return self.name == value

# 기존에는 name 을 통해 변수의 이름을 가져올 수 있었다.
>>> Fruits.APPLE.name == "APPLE"
True

# __eq__ 함수를 통해 "==" 연산시 일어나는 결과를 변경할 수 있다.
# "==" 비교시 name 을 추가로 기입하지 않아도 된다.
# 만약, __eq__ 함수의 로직을 추가해주지 않았다면, 아래 결과는 False가 뜬다.
>>> Fruits.APPLE == "APPLE"
True

Enum 을 쓰다보면, 열거형 멤버의 이름(name)을 사용하여 값을 비교하는 경우가 생길 수 있다.

만약, Fruits 의 APPLE 의 이름을 문자열 "APPLE" 과 비교한다면, Fruits.APPLE.name == "APPLE" 로 써야하는 불편함이 있다.

이때, 필요연산에 따라 적절한 함수(여기서는 __eq__)를 정의 해주면, Fruits.APPLE == "APPLE" 로 연산 사용할 수 있다.