Автоматическое присваивание значений в Enum

Представим ситуацию: вы любите все структурировать и храните много данных в енамах: Enum, StrEnum, IntEnum и так далее.

В какой-то момент обязательно накопится большое количество структур, где символическое имя будет равно значению, как в примере ниже:

from enum import Enum  
  
  
class MyMood(Enum):  
    HAPPY = "HAPPY"  
    SAD = "SAD"

Как красиво выйти из этой ситуации? Об этом мы и поговорим.

Функция auto()

В первую очередь, стоит обратить внимание на класс auto из Enum.

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

from enum import Enum, auto  
  
  
class MyMood(Enum):  
    HAPPY = auto()  
    SAD = auto()

Как это работает

Класс auto автоматически присваивает значение для символического имени в виде номера по порядку.
Например, для блока кода выше будут присвоены значения HAPPY = 1, SAD = 2:

print(MyMood.HAPPY.value, MyMood.SAD.value)
> 1 2

Когда это может пригодиться?

Например, когда мы перебираем Enum-ы через match-case и значения нам не важны.

А если важны?

Что делать, если в value нужно хранить именно символическое имя?

Представим ситуацию, когда нам очень нужно, чтобы name было равно value, и чтоб это value присваивалось автоматически.

Здесь нам на помощь придет переопределение метода _generate_next_value_.

Синтаксис

_generate_next_value_(name, start, count, last_values)
Этот метод и отвечает за увеличение значения символического имени на 1 в исходной реализации, но мы можем подстроить его под себя.

Это можно сделать прямо внутри нужного нам Enum-а:

class MyMood(Enum):  
    @staticmethod  
    def _generate_next_value_(name, start, count, last_values):  
        return name  
  
    HAPPY = auto()  
    SAD = auto()  

главное, переопределить метод перед членами Enum-а, иначе словим исключение `TypeError.

И тогда при выводе значений в терминал мы получим уже имена членов Enum, а не индексы:

print(MyMood.HAPPY.value, MyMood.SAD.value)
> HAPPY SAD

Если ситуация часто повторяется, то можно сделать одно общее решение для всех. Создадим класс AutoName, унаследованный от Enum и переопределим метод _generate_next_value_:

from enum import Enum, auto  


class AutoName(Enum):  
    @staticmethod  
    def _generate_next_value_(name, start, count, last_values):  
        return name  

Нам остается отнаследоваться от этого класса и использовать auto() в качестве значений символических имен Enum-а:

class MyMood(AutoName):  
    HAPPY = auto()  
    SAD = auto()

Убедимся, что наша шалость удалась:

print(MyMood.HAPPY.value, MyMood.SAD.value)
> HAPPY SAD

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