Python:enum
lastmod 2024-07-29created 2021-03-129 edits
An enumeration is a set of symbolic names (members) bound to unique constant values. Within an enumeration, the members can be compared by identity, and the enumeration itself can be iterated over.
Since version 3.4, Python has supported the enum (enumeration) type like other languages. An enum is generally used to define a set of related constants. Because an enum class lets you restrict the kinds of instances, it helps you write robust programs.
List of member names
from enum import Enum
from typing import List
class Command(Enum):
unknown = 0
core = 1
task = 2
ctrl = 3
daemon = 4
def get_available_commands() -> List[str]:
return [n for n, _ in Command.__members__.items() if n != Command.unknown.name]
Class-based type definition
As a simple example, let's write an enum class representing three key web technologies (HTML, CSS, JS). We create the Skill type by extending the Enum class imported from the built-in enum module.
A constant instance of an enum type has a name (name) and a value (value) property by default.
Enum types are iterable, so you can easily check all constants with a for statement.
Programmatic access to enumeration members and their attributes
Sometimes it's useful to access members in enumerations programmatically (i.e. situations where Color.RED won't do because the exact color is not known at program-writing time). Enum allows such access:
If you want to access enum members by name, use item access:
If you have an enum member and need its name or value:
Functional type definition
It isn't used often, but instead of extending the Enum class you can also define an enum type by calling it like a regular function.
>>> Skill = Enum("Skill", "HTML CSS JS")
>>> list(Skill)
[<Skill.HTML: 1>, <Skill.CSS: 2>, <Skill.JS: 3>]
Automatic value assignment
When using enum, the actual value (value) often doesn't matter much. In that case, the auto() helper function of the enum module assigns a unique number to every constant, incrementing by 1: 1 to the first constant, 2 to the second, and so on.
Using the auto() function also has the advantage that you can add new constants without worrying about which numbers were assigned to the existing ones.
In addition, by overriding the _generate_next_value_() method of the Enum class, you can automatically assign values other than numbers. For example, you can automatically assign a string identical to the constant's name as its value.
from enum import Enum, auto
class Skill(Enum):
def _generate_next_value_(name, start, count, last_values):
return name
HTML = auto()
CSS = auto()
JS = auto()
enum mixin
One possible inconvenience of enum types is that you have to use the name or value property to access a constant's name or value.
>>> Skill.HTML.name == 'HTML'
True
>>> Skill.HTML.value == 1
True
>>> Skill.HTML == 'HTML'
False
>>> type(Skill.HTML)
<enum 'Skill'>
This is because every constant is ultimately an instance of that enum type. But when coding with enum types, using name or value every time is tedious and easy to forget.
In that case, use the enum mixin technique to write an enum class that extends str.
class StrEnum(str, Enum):
def _generate_next_value_(name, start, count, last_values):
return name
def __repr__(self):
return self.name
def __str__(self):
return self.name
Then define enum classes by extending this class.
Now the constants in the Skill type are treated fully as strings, which makes them more convenient to use.
Extending Enum
Like any other regular class, the Enum class can be extended and used in many ways.
from enum import Enum
class Skill(Enum):
HTML = ("HTML", "Hypertext Markup Language")
CSS = ("CSS", "Cascading Style Sheets")
JS = ("JS", "JavaScript")
def __init__(self, title, description):
self.title = title
self.description = description
@classmethod
def get_most_popular(cls):
return cls.JS
def lower_title(self):
return self.title.lower()
>>> Skill.HTML.value
('HTML', 'Hypertext Markup Language')
>>> Skill.HTML.title
'HTML'
>>> Skill.HTML.description
'Hypertext Markup Language'
>>> Skill.get_most_popular()
<Skill.JS: ('JS', 'JavaScript')>
>>> Skill.CSS.lower_title()
'css'
StrEnum
If you want to use StrEnum, available since Python 3.11, in earlier versions too:
# -*- coding: utf-8 -*-
import sys
from enum import Enum
class _StrEnum(str, Enum):
def __new__(cls, *values):
if len(values) > 3:
raise TypeError("too many arguments for str(): %r" % (values,))
if len(values) == 1:
# it must be a string
if not isinstance(values[0], str):
raise TypeError("%r is not a string" % (values[0],))
if len(values) >= 2:
# check that encoding argument is a string
if not isinstance(values[1], str):
raise TypeError("encoding must be a string, not %r" % (values[1],))
if len(values) == 3:
# check that errors argument is a string
if not isinstance(values[2], str):
raise TypeError("errors must be a string, not %r" % (values[2]))
value = str(*values)
member = str.__new__(cls, value)
member._value_ = value
return member
# noinspection PyMethodParameters
def _generate_next_value_(name, start, count, last_values):
return name.lower()
if sys.version_info[:2] >= (3, 11):
from enum import StrEnum # noqa
else:
StrEnum = _StrEnum
Utilities
Here are a few utility functions.
# -*- coding: utf-8 -*-
from enum import Enum
from typing import Dict, Type, TypeVar
_ASSERT_MSG_NOT_ENUM_FORMAT = "The `{cls}` type must inherit from Enum class."
_T = TypeVar("_T")
Given the above...
Converting a string to an enum
def string_to_enum(name: str, cls: Type[_T]) -> _T:
assert issubclass(cls, Enum), _ASSERT_MSG_NOT_ENUM_FORMAT.format(cls=cls.__name__)
enums = [s for s in dir(cls) if not s.startswith("_")]
if name in enums:
return getattr(cls, name)
raise KeyError(f"Not found `{name}` enum in {cls.__name__}")
Returning a dictionary mapping string keys to enum values
def string_to_enum_map(cls: Type[_T]) -> Dict[str, _T]:
assert issubclass(cls, Enum), _ASSERT_MSG_NOT_ENUM_FORMAT.format(cls=cls.__name__)
enums = [s for s in dir(cls) if not s.startswith("_")]
return {e: getattr(cls, e) for e in enums}
Returning a dictionary mapping string (lowercase) keys to enum values
def lower_string_to_enum_map(cls: Type[_T]) -> Dict[str, _T]:
assert issubclass(cls, Enum), _ASSERT_MSG_NOT_ENUM_FORMAT.format(cls=cls.__name__)
enums = [s for s in dir(cls) if not s.startswith("_")]
return {e.lower(): getattr(cls, e) for e in enums}
See also
- Python
- enum