Гайд по использованию enum в Python

Модуль enum содержит в себе тип для перечисления значений с возможностью итерирования и сравнения. Его можно использовать для создания понятных обозначений вместо использования чисел (для которых приходится помнить, какое число что обозначает) или строк (в которых легко опечататься и не заметить).

Создание

Для создания перечисления необходимо создать класc, являющийся наследником класса enum.Enum. Для установки значений нужно добавить соответствующие атрибуты в класс. Пример использования:

#  enum_create.py
import enum

class BugStatus(enum.Enum):

    new = 7
    incomplete = 6
    invalid = 5
    wont_fix = 4
    in_progress = 3
    fix_committed = 2
    fix_released = 1

print('\nMember name: {}'.format(BugStatus.wont_fix.name)) 
print('Member value: {}'.format(BugStatus.wont_fix.value))

Атрибуты класса Enum конвертируются в экземпляры при парсинге. Каждый экземпляр имеет параметр name, в котором хранится название, а также value, в котором хранится установленное значение.

$ python3 enum_create.py

Member name: wont_fix
Member value: 4

Итерирование

При итерировании по классу вы пройдёте по атрибутам.

#  enum_iterate.py
import enum

class BugStatus(enum.Enum):

    new = 7
    incomplete = 6
    invalid = 5
    wont_fix = 4
    in_progress = 3
    fix_committed = 2
    fix_released = 1

for status in BugStatus:
    print('{:15} = {}'.format(status.name, status.value))

Цикл будет идти по элементам в том порядке, в котором они указаны при создании класса. Названия и значения никак не влияют на порядок.

$ python3 enum_iterate.py

new             = 7
incomplete      = 6
invalid         = 5
wont_fix        = 4
in_progress     = 3
fix_committed   = 2
fix_released    = 1

Сравнение перечислений

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

#  enum_comparison.py
import enum

class BugStatus(enum.Enum):

    new = 7
    incomplete = 6
    invalid = 5
    wont_fix = 4
    in_progress = 3
    fix_committed = 2
    fix_released = 1

actual_state = BugStatus.wont_fix
desired_state = BugStatus.fix_released

print('Equality:',
      actual_state == desired_state,
      actual_state == BugStatus.wont_fix)  # проверка на равенство
print('Identity:',
      actual_state is desired_state,
      actual_state is BugStatus.wont_fix)  # проверка на то, это один и тот же элемент или нет
print('Ordered by value:')
try:
    print('\n'.join('  ' + s.name for s in sorted(BugStatus)))  # пытаемся упорядочить
except TypeError as err:
    print('  Cannot sort: {}'.format(err))  # вывод ошибки в случае неудачи

Операторы больше и меньше порождают TypeError.

$ python3 enum_comparison.py

Equality: False True
Identity: False True
Ordered by value:
  Cannot sort: '<' not supported between instances of 'BugStatus' and 'BugStatus'

IntEnum

Для создания перечислений с возможностью сравнения можно использовать IntEnum, который поддерживает сравнение по значениям.

#  enum_intenum.py
import enum

class BugStatus(enum.IntEnum):

    new = 7
    incomplete = 6
    invalid = 5
    wont_fix = 4
    in_progress = 3
    fix_committed = 2
    fix_released = 1

print('Ordered by value:')
print('\n'.join('  ' + s.name for s in sorted(BugStatus)))  # упорядочивание по значению

$ python3 enum_intenum.py

Ordered by value:
  fix_released
  fix_committed
  in_progress
  wont_fix
  invalid
  incomplete
  new

Уникальные значения в перечислениях

Элементы перечисления с одинаковыми значениями являются несколькими названиями, указывающими на один и тот же объект.

#  enum_aliases.py
import enum

class BugStatus(enum.Enum):

    new = 7
    incomplete = 6
    invalid = 5
    wont_fix = 4
    in_progress = 3
    fix_committed = 2
    fix_released = 1

    by_design = 4
    closed = 1

for status in BugStatus:
    print('{:15} = {}'.format(status.name, status.value))

print('\nSame: by_design is wont_fix: ',
      BugStatus.by_design is BugStatus.wont_fix)
print('Same: closed is fix_released: ',
      BugStatus.closed is BugStatus.fix_released)

Так как by_design и closed являются синонимами для других элементов, то они не появляются как элементы в циклах. Истинным считается название, указанное первым при объявлении.

$ python3 enum_aliases.py

new             = 7
incomplete      = 6
invalid         = 5
wont_fix        = 4
in_progress     = 3
fix_committed   = 2
fix_released    = 1

Same: by_design is wont_fix:  True
Same: closed is fix_released:  True

Если вы хотите, чтобы все элементы обязательно имели разные значения, то добавьте декоратор @unique перед объявлением класса.

#  enum_unique_enforce.py
import enum

@enum.unique
class BugStatus(enum.Enum):

    new = 7
    incomplete = 6
    invalid = 5
    wont_fix = 4
    in_progress = 3
    fix_committed = 2
    fix_released = 1

    # This will trigger an error with unique applied.
    by_design = 4
    closed = 1

Элементы с повторяющимися значениями будут вызывать ValueError во время интерпретации.

$ python3 enum_unique_enforce.py

Traceback (most recent call last):
  File "enum_unique_enforce.py", line 11, in <module>
    class BugStatus(enum.Enum):
  File ".../lib/python3.7/enum.py", line 848, in unique
    (enumeration, alias_details))
ValueError: duplicate values found in <enum 'BugStatus'>:
by_design -> wont_fix, closed -> fix_released

Другой способ создания перечислений

Иногда удобнее не хардкодить элементы перечисления, а указывать их в более удобном виде. Для этого можно передать значения в конструктор класса:

#  enum_programmatic_create.py
import enum

BugStatus = enum.Enum(
    value='BugStatus',
    names=('fix_released fix_committed in_progress '
           'wont_fix invalid incomplete new'),
)

print('Member: {}'.format(BugStatus.new))

print('\nAll members:')
for status in BugStatus:
    print('{:15} = {}'.format(status.name, status.value))

Аргумент value является названием перечисления, которое используется для создания представления элементов. Второй аргумент names принимает список названий в перечислении. Если подать одну строку, то она будет разбита по пробельным символам и запятым, а значения будут числами, начиная с 1.

$ python3 enum_programmatic_create.py

Member: BugStatus.new

All members:
fix_released    = 1
fix_committed   = 2
in_progress     = 3
wont_fix        = 4
invalid         = 5
incomplete      = 6
new             = 7

Также аргумент name принимает список пар название — значение, либо аналогичный словарь.

#  enum_programmatic_mapping.py
import enum

BugStatus = enum.Enum(
    value='BugStatus',
    names=[
        ('new', 7),
        ('incomplete', 6),
        ('invalid', 5),
        ('wont_fix', 4),
        ('in_progress', 3),
        ('fix_committed', 2),
        ('fix_released', 1),
    ],
)

print('All members:')
for status in BugStatus:
    print('{:15} = {}'.format(status.name, status.value))

Передача списка пар позволяет сохранить порядок элементов, аналогично случаю, где мы объявляли атрибуты.

$ python3 enum_programmatic_mapping.py

All members:
new             = 7
incomplete      = 6
invalid         = 5
wont_fix        = 4
in_progress     = 3
fix_committed   = 2
fix_released    = 1

Другие типы значений

Значения элементов перечислений необязательно должны быть числами, они могут иметь любой тип. Если значение имеет тип tuple, то элементы передаются как отдельные аргументы в функцию __init()__.

#  enum_tuple_values.py
import enum

class BugStatus(enum.Enum):

    new = (7, ['incomplete',
               'invalid',
               'wont_fix',
               'in_progress'])
    incomplete = (6, ['new', 'wont_fix'])
    invalid = (5, ['new'])
    wont_fix = (4, ['new'])
    in_progress = (3, ['new', 'fix_committed'])
    fix_committed = (2, ['in_progress', 'fix_released'])
    fix_released = (1, ['new'])

    def __init__(self, num, transitions):
        self.num = num
        self.transitions = transitions

    def can_transition(self, new_state):
        return new_state.name in self.transitions

print('Name:', BugStatus.in_progress)
print('Value:', BugStatus.in_progress.value)
print('Custom attribute:', BugStatus.in_progress.transitions)
print('Using attribute:',
      BugStatus.in_progress.can_transition(BugStatus.new))

В этом примере каждое значение является парой из числового id и списка строк, описывающих возможный переход из данного состояния.

$ python3 enum_tuple_values.py

Name: BugStatus.in_progress
Value: (3, ['new', 'fix_committed'])
Custom attribute: ['new', 'fix_committed']
Using attribute: True

Для более сложных случаев tuple может быть плохим решением. Так как типом value может быть что угодно, мы можем использовать словари, чтобы обозначить различные параметры. Сложные значения передаются прямиком в __init()__ как единственный аргумент помимо self.

#  enum_complex_values.py
import enum

class BugStatus(enum.Enum):

    new = {
        'num': 7,
        'transitions': [
            'incomplete',
            'invalid',
            'wont_fix',
            'in_progress',
        ],
    }
    incomplete = {
        'num': 6,
        'transitions': ['new', 'wont_fix'],
    }
    invalid = {
        'num': 5,
        'transitions': ['new'],
    }
    wont_fix = {
        'num': 4,
        'transitions': ['new'],
    }
    in_progress = {
        'num': 3,
        'transitions': ['new', 'fix_committed'],
    }
    fix_committed = {
        'num': 2,
        'transitions': ['in_progress', 'fix_released'],
    }
    fix_released = {
        'num': 1,
        'transitions': ['new'],
    }

    def __init__(self, vals):
        self.num = vals['num']
        self.transitions = vals['transitions']

    def can_transition(self, new_state):
        return new_state.name in self.transitions

print('Name:', BugStatus.in_progress)
print('Value:', BugStatus.in_progress.value)
print('Custom attribute:', BugStatus.in_progress.transitions)
print('Using attribute:',
      BugStatus.in_progress.can_transition(BugStatus.new))

Данный пример аналогичен предыдущему, но в нём используются словари вместо tuple для удобства.

$ python3 enum_complex_values.py

Name: BugStatus.in_progress
Value: {'num': 3, 'transitions': ['new', 'fix_committed']}
Custom attribute: ['new', 'fix_committed']
Using attribute: True