Автоматическое присваивание значений в 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
Мы получили то, что нам было нужно, избавились от дублирования и лишней копипасты, сделав код еще немного чище.