Customize and Extend Python Enum Class

출처 :https://www.pythontutorial.net/python-oop/python-enum-class/

Customize and Extend Python Enum Class

 

 

Customize Python enum classes

 

파이썬 열거형은 클래스입니다. 즉, 메서드를 추가하거나 dunder 메서드를 구현하여 동작을 사용자 지정할 수 있습니다.

다음 예제에서는 PaymentStatus 열거형 클래스를 정의합니다.

 

from enum import Enum


class PaymentStatus(Enum):
    PENDING = 1
    COMPLETED = 2
    REFUNDED = 3

 

PaymentStatus 열거형에는 세 가지 멤버가 있습니다 : PENDING, COMPLETED, and REFUND.

다음은 PaymentStatus의 멤버를 표시합니다.

 

print(PaymentStatus.PENDING)  # PaymentStatus.PENDING

 

PaymentStatus 멤버가 문자열에 표시되는  방식을 사용자 지정하려면 __str__ 메서드를 구현할 수 있습니다  . 예를 들어:

 

from enum import Enum


class PaymentStatus(Enum):
    PENDING = 1
    COMPLETED = 2
    REFUNDED = 3

    def __str__(self):
        return f'{self.name.lower()}({self.value})'


print(PaymentStatus.PENDING)  # pending(1)

 

Implementing __eq__ method

 

다음은 PaymentStatus 열거형 클래스의 멤버를 정수와 비교하려고 시도합니다.

 

if PaymentStatus.PENDING == 1:
    print('The payment is pending.')

 

PaymentStatus.PENDING이 정수 1과 같지 않기 때문에 아무 것도 표시되지 않습니다.

PaymentStatus 멤버와 정수 간의 비교를 허용하려면  다음과 같이 __eq__ 메서드를 구현할 수 있습니다.

 

from enum import Enum


class PaymentStatus(Enum):
    PENDING = 1
    COMPLETED = 2
    REFUNDED = 3

    def __str__(self):
        return f'{self.name.lower()}({self.value})'

    def __eq__(self, other):
        if isinstance(other, int):
            return self.value == other

        if isinstance(other, PaymentStatus):
            return self is other

        return False


if PaymentStatus.PENDING == 1:
    print('The payment is pending.')

 

In the __eq__ method:

• 비교할 값이 정수이면 멤버의 값을 정수와 비교합니다.
• 비교할 값이 PaymentStatus 열거형의 인스턴스인 경우 is 연산자를 사용하여 값을 PaymentStatus 멤버의 멤버와 비교합니다.
• 그렇지 않으면 거짓을반환합니다.
  

 

output:

 

The payment is pending.

 

Implementing __lt__ method

 

PaymentStatus의 멤버가  해당 값에 따라 정렬 순서를 따르도록 하려는 경우를 가정해 보겠습니다. 또한 정수와 비교하려고합니다.

이렇게 하려면 __lt__ 메서드를 구현하고 functools 모듈의 @total_ordering 데코레이터를  사용할 수 있습니다.

 

from enum import Enum
from functools import total_ordering


@total_ordering
class PaymentStatus(Enum):
    PENDING = 1
    COMPLETED = 2
    REFUNDED = 3

    def __str__(self):
        return f'{self.name.lower()}({self.value})'

    def __eq__(self, other):
        if isinstance(other, int):
            return self.value == other

        if isinstance(other, PaymentStatus):
            return self is other

        return False

    def __lt__(self, other):
        if isinstance(other, int):
            return self.value < other

        if isinstance(other, PaymentStatus):
            return self.value < other.value

        return False


# compare with an integer
status = 1
if status < PaymentStatus.COMPLETED:
    print('The payment has not completed')

# compare with another member
status = PaymentStatus.PENDING
if status >= PaymentStatus.COMPLETED:
    print('The payment is not pending')

 

Implementing the __bool__ method

 

기본적으로 열거형의 모든 멤버는 true입니다. 예를 들어:

 

for member in PaymentStatus:
    print(member, bool(member))

 

이 동작을 사용자 지정하려면 __bool__ 메서드를 구현해야 합니다  . COMPLETED는 TRUE, REFUNDED 및 PENDING 는 FALSE로  설정한다고 가정합니다.

다음은 이 논리를 구현하는 방법을 보여 줍니다.

 

from enum import Enum
from functools import total_ordering


@total_ordering
class PaymentStatus(Enum):
    PENDING = 1
    COMPLETED = 2
    REFUNDED = 3

    def __str__(self):
        return f'{self.name.lower()}({self.value})'

    def __eq__(self, other):
        if isinstance(other, int):
            return self.value == other

        if isinstance(other, PaymentStatus):
            return self is other

        return False

    def __lt__(self, other):
        if isinstance(other, int):
            return self.value < other

        if isinstance(other, PaymentStatus):
            return self.value < other.value

        return False

    def __bool__(self):
        if self is self.COMPLETED:
            return True

        return False


for member in PaymentStatus:
    print(member, bool(member))

 

output:

 

pending(1) False
completed(2) True
refunded(3) False

 

Extend Python enum classes

 

파이썬은 멤버가 있으면 열거형 클래스를 확장 할 수 없습니다 . 그러나 이것은 제한이 아닙니다. 메서드는 있지만 멤버는 없는 base 클래스를 정의한 다음 이 기본 클래스를 확장할 수 있기 때문입니다. 예를 들어:

 

먼저  values에 따라 멤버를 정렬하는 OrderedEnum base 클래스를 정의합니다.

 

from enum import Enum
from functools import total_ordering


@total_ordering
class OrderedEnum(Enum):
    def __lt__(self, other):
        if isinstance(other, OrderedEnum):
            return self.value < other.value
        return NotImplemented

 

둘째, OrderedEnum클래스를 확장하는 ApprovalStatus 정의.

 

class ApprovalStatus(OrderedEnum):
    PENDING = 1
    IN_PROGRESS = 2
    APPROVED = 3

 

셋째, ApprovalStatus열거형 클래스의 멤버를 비교합니다.

 

status = ApprovalStatus(2)
if status < ApprovalStatus.APPROVED:
    print('The request has not been approved.')

Output:

The request has not been approved.

 

전체 Code

 

from enum import Enum
from functools import total_ordering


@total_ordering
class OrderedEnum(Enum):
    def __lt__(self, other):
        if isinstance(other, OrderedEnum):
            return self.value < other.value
        return NotImplemented


class ApprovalStatus(OrderedEnum):
    PENDING = 1
    IN_PROGRESS = 2
    APPROVED = 3


status = ApprovalStatus(2)
if status < ApprovalStatus.APPROVED:
    print('The request has not been approved.')

 

Summary

 

• dunder 메서드를 구현하여 Python 열거형 클래스의 동작을 사용자 지정합니다.
• 멤버가 없고와 메서드가 있는 emum 클래스를 정의하고 이 base 클래스를, 멤버를 추가하는 클래스를 상속시켜 열거형을 확장합니다.