Настройка и расширение класса перечисления в Python

В этом руководстве вы узнаете, как настраивать и расширять пользовательские классы перечислений в Python.

Настройка классов перечислений в Python

Перечисления Python — это классы. Это означает, что вы можете добавлять к ним методы или реализовывать методы dunder для настройки их поведения.

В следующем примере определяется класс перечисления PaymentStatus:

from enum import Enum


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

Перечисление PaymentStatus состоит из трех членов: PENDING, COMPLETED и 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)

Реализация метода __eq__

В следующем примере член класса перечисления 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.')

В методе __eq__:

• Если сравниваемое значение является целым числом, оно сравнивает значение элемента с целым числом.
• Если сравниваемое значение является экземпляром перечисления PaymentStatus, оно сравнивает значение с элементом элемента PaymentStatus с помощью оператора is.
• В противном случае он возвращает False.

Программа работает как положено и возвращает следующий результат:

The payment is pending.

Реализация метода __lt__

Предположим, вы хотите, чтобы члены PaymentStatus имели порядок сортировки на основе их значения. И вы также хотите сравнить их с целым числом.

Для этого вы можете реализовать метод __lt__ и использовать декоратор @total_ordering из модуля functools:

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')

Реализация метода __bool__

По умолчанию все члены перечисления правдивы. Например:

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

Чтобы настроить это поведение, вам необходимо реализовать метод __bool__. Предположим, вы хотите, чтобы элементы COMPLETED и REFUNDED имели значение True, а 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))

Программа вывела следующее:

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

Расширение класса перечислений в Python

Python не позволяет расширять класс перечисления, если у него нет элементов. Однако это не ограничение, потому что вы можете определить базовый класс, у которого есть методы, но нет членов, а затем расширить этот базовый класс.

Пример:

• Сначала определите базовый класс OrderedEnum, который упорядочивает члены по их значениям:

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

• Во-вторых, определите ApprovalStatus, который расширяет класс OrderedEnum:

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.')

Выход:

The request has not been approved.

Положим все это вместе:

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.')