[Python] Enum 타입

파이썬 3.4부터 enum 타입을 지원하고 있습니다. 보통 enum은 연관 있는 여러 개의 상수를 선언할 때 사용하게 됩니다. enum 클래스를 사용하면 인스턴스의 종류를 제한할 수 있어서 프로그래밍에 있어 도움이 됩니다.

클래스 타입

enum 타입의 상수는 기본적으로 name과 value를 가지고 있습니다.

from enum import Enum


class Status(Enum):
    SUCCESS = 1
    FAILURE = 2
    RETRY = 3


Status.SUCCESS
Status.SUCCESS.name
Status.SUCCESS.value

# Status.SUCCESS
# SUCCESS
# 1

이터레이션을 지원하므로 반복문을 통해 접근할 수 있고 value나 name으로도 직접 접근이 가능합니다.

for status in Status:
    print(status)


print(Status(1))
print(Status['FAILURE'])


# Status.SUCCESS
# Status.FAILURE
# Status.RETRY


# Status.SUCCESS
# Status.FAILURE

또한 상수의 name은 중복이 불가능하지만 value는 중복이 가능합니다. value가 중복일 경우, 나중에 선언된 name은 먼저 선언된 name의 alias로 동작하게 됩니다. 

from enum import Enum


class Status(Enum):
    SUCCESS = 1
    FAILURE = 2
    RETRY = 3
    RETRY_FOR_3_WAYS = 3


print(Status.RETRY)
print(Status.RETRY_FOR_3_WAYS)
print(Status(3))


# Status.RETRY
# Status.RETRY
# Status.RETRY

만약 value의 unique를 보장하려면 아래와 같이 데코레이터를 사용해야 합니다.

from enum import Enum, unique

@unique
class Status(Enum):
    SUCCESS = 1
    FAILURE = 2
    RETRY = 3
    RETRY_FOR_3_WAYS = 3 # 에러

값 자동 할당하기

만약 value가 중요하지 않으면 자동으로 값을 할당하게 할 수 있습니다.

from enum import Enum, auto


class Status(Enum):
    SUCCESS = auto()
    FAILURE = auto()
    RETRY = auto()


list(Status)


# [<Status.SUCCESS: 1>, <Status.FAILURE: 2>, <Status.RETRY: 3>]

값을 자동으로 할당받게 하면 어떠한 값을 선언했는지 신경쓰지 않고 설정할 수 있습니다. 위의 auto()의 값은 _generate_next_value() 메소드에 의해 결정됩니다. 만약 해당 메소드로 값을 수정하려면 가장 먼저 정의가 되어야 합니다.

from enum import Enum, auto


class Status(Enum):
    def _generate_next_value_(name, start, count, last_values):
        return name

    SUCCESS = auto()
    FAILURE = auto()
    RETRY = auto()


list(Status)

# [<Status.SUCCESS: 'SUCCESS'>, <Status.FAILURE: 'FAILURE'>, <Status.RETRY: 'RETRY'>]


### 


class Status(Enum):
    SUCCESS = auto()
    FAILURE = auto()
    RETRY = auto()

    def _generate_next_value_(name, start, count, last_values): # 가장 먼저 재정의가 되지 않아서 오류!
        return name

mixin으로 간편한 enum 만들기

enum 클래스로 정의하고 이름을 선언하려면 name이나 value 속성을 매번 선언해야 하는데 위 _generate_next_value() 를 사용하여 간편하게 변경할 수 있습니다. 아래는 mixin 기법을 사용해 str 클래스를 확장하는 예제입니다.

from enum import Enum, auto


class StrEnum(str, Enum):
    def _generate_next_value_(name, start, count, last_values):
        return name

    def __str__(self):
        return self.name


class Status(StrEnum):
    SUCCESS = auto()
    FAILURE = auto()
    RETRY = auto()


print(Status.SUCCESS)
print(Status.SUCCESS == 'SUCCESS')
print(isinstance(Status.SUCCESS, str))
print(type(Status.SUCCESS))


# SUCCESS
# True
# True
# <enum 'Status'>

enum 클래스 확장

enum 클래스를 확장하여 설명이나 메소드를 정의하고 이를 사용할 수 있습니다.

class Status(Enum):
    SUCCESS = (1, 'success')
    FAILURE = (2, 'failure')
    RETRY = (3, 'retry')

    def __init__(self, title, description):
        self.title = title
        self.description = description

    @classmethod
    def most(cls):
        return cls.SUCCESS

    def upper_description(self):
        return self.description.upper()


print(Status.SUCCESS)
print(Status.SUCCESS.value)
print(Status.SUCCESS.title)
print(Status.SUCCESS.description)


# Status.SUCCESS
# (1, 'success')
# 1
# success


print(Status.most())
print(Status.SUCCESS.upper_description())


# Status.SUCCESS
# SUCCESS

IntEnum

위에서 mixin 기법을 활용해 StrEnum을 만들었지만 정수형에 해당해선 enum 클래스에서 지원을 해줍니다.

from enum import Enum, IntEnum


class Status(IntEnum):
    SUCCESS = 1
    FAILURE = 2
    RETRY = 3


class Status1(IntEnum):
    SUCCESS1 = 1
    FAILURE1 = 2
    RETRY1 = 3


print(Status.SUCCESS)
print(Status.SUCCESS == 1)
print(Status.SUCCESS == Status1.SUCCESS1)

# Status.SUCCESS
# True
# True

a = [0, 1, 2, 3, 4, 5]
print(a[Status.SUCCESS])

# 1


class EnumStatus(Enum):
    SUCCESS1 = 1
    FAILURE1 = 2
    RETRY1 = 3

print(EnumStatus.SUCCESS1)

# EnumStatus.SUCCESS1