Personalize e estenda a classe Python Enum

Resumo : neste tutorial, você aprenderá como personalizar e estender as classes enum personalizadas do Python.

Personalize classes enum do Python

As enumerações do Python são classes . Isso significa que você pode adicionar métodos a eles ou implementar métodos dunder para personalizar seus comportamentos.

O exemplo a seguir define a PaymentStatusclasse de enumeração:

from enum import Enum


class PaymentStatus(Enum):
    PENDING = 1
    COMPLETED = 2
    REFUNDED = 3Linguagem de código:  Python  ( python )

A PaymentStatusenumeração tem três membros: PENDING, COMPLETED, e REFUND.

O seguinte exibe o membro do PaymentStatusmembro ‘:

print(PaymentStatus.PENDING)Linguagem de código:  Python  ( python )

Ele mostra o seguinte:

PaymentStatus.PENDINGLinguagem de código:  Python  ( python )

Para personalizar como o PaymentStatusmembro é representado na string, você pode implementar o __str__método. Por exemplo:

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)Linguagem de código:  Python  ( python )

Agora, ele retorna a seguinte string:

pending(1)Linguagem de código:  Python  ( python )

Implementando o método __eq__

As seguintes tentativas de comparar um membro da PaymentStatusclasse enum com um número inteiro:

if PaymentStatus.PENDING == 1:
    print('The payment is pending.')Linguagem de código:  Python  ( python )

Não mostra nada porque the PaymentStatus.PENDINGnão é igual ao número inteiro 1.

Para permitir a comparação entre PaymentStatusmembro e um número inteiro, você pode implementar o __eq__método assim:

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.')Linguagem de código:  Python  ( python )

No __eq__método:

• Se o valor a ser comparado for um número inteiro, ele compara o valor do membro com o número inteiro.
• Se o valor a ser comparado for uma instância da PaymentStatusenumeração, ele compara o valor com o membro do PaymentStatusmembro usando o isoperador.
• Caso contrário, ele retorna False.

O programa funciona conforme o esperado e retorna a seguinte saída:

The payment is pending.Linguagem de código:  Python  ( python )

Implementando o método __lt__

Suponha que você queira que os membros a PaymentStatusseguir tenham uma ordem de classificação baseada em seus valores. E você também deseja compará-los com um número inteiro.

Para fazer isso, você pode implementar o __lt__método e usar o @total_orderingdecorador do functoolsmódulo:

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')Linguagem de código:  Python  ( python )

Implementando o método __bool__

Por padrão, todos os membros de uma enumeração são verdadeiros. Por exemplo:

for member in PaymentStatus:
    print(member, bool(member))Linguagem de código:  Python  ( python )

Para personalizar esse comportamento, você precisa implementar o __bool__método. Suponha que você queira que os membros COMPLETEDand REFUNDEDsejam True enquanto os PENDINGto be False.

A seguir mostramos como implementar essa lógica:

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))Linguagem de código:  Python  ( python )

O programa produz o seguinte:

pending(1) False
completed(2) True
refunded(3) FalseLinguagem de código:  Python  ( python )

Estenda as classes enum do Python

Python não permite estender uma classe enum, a menos que ela não tenha nenhum membro. No entanto, isso não é uma limitação. Porque você pode definir uma classe base que possui métodos, mas nenhum membro e então estender essa classe base. Por exemplo:

Primeiro, defina a OrderedEnumclasse base que ordena os membros pelos seus valores:

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 NotImplementedLinguagem de código:  Python  ( python )

Segundo, defina o ApprovalStatusque estende a OrderedEnumclasse:

class ApprovalStatus(OrderedEnum):
    PENDING = 1
    IN_PROGRESS = 2
    APPROVED = 3Linguagem de código:  Python  ( python )

Terceiro, compare os membros da ApprovalStatusclasse enum:

status = ApprovalStatus(2)
if status < ApprovalStatus.APPROVED:
    print('The request has not been approved.')Linguagem de código:  Python  ( python )

Saída:

The request has not been approved.Linguagem de código:  Python  ( python )

Junte tudo:

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.')Linguagem de código:  Python  ( python )

Resumo

• Implemente métodos dunder para personalizar o comportamento das classes enum do Python.
• Defina uma classe emum sem membros e métodos e estenda esta classe base.