How to create an indexable and sliceable enum with string values in Python?

Question

I've got a file like this:

class Level(Enum):
    prerequisite_level: Optional["Level"]
    dependent_level: Optional["Level"]
    lower_priority_levels: List["Level"]
    greater_priority_levels: List["Level"]

    DATA_CHECK = "data check"
    DESIGN_CHECK = "design check"
    ALERT = "alert"

The enum values are in a specific order, and based on each of those levels I need to be able to get the previous one, the next one, and all the previous and next ones. I believe I need to be able to index the levels numerically to get these values, so I've added a constant to be able to do this:

INCREASING_PRIORITY_LEVELS: List[Level] = list(Level)

for priority_level_index, threshold_level in enumerate(Level):
    if priority_level_index > 0:
        threshold_level.prerequisite_level = Level[priority_level_index - 1]
    else:
        threshold_level.prerequisite_level = None

    if priority_level_index < len(Level) - 1:
        threshold_level.dependent_level = Level[priority_level_index + 1]
    else:
        threshold_level.dependent_level = None

    threshold_level.lower_priority_levels = Level[:priority_level_index]
    threshold_level.greater_priority_levels = Level[priority_level_index + 1:]

This is clunky, and I'd like to get rid of this constant. Do I need to implement __getitem__ or something to make this possible?

Answer 1

Another version building on @blhsing's answer with types and support for finding the index of a member. Unfortunately mypy does not fully support the metaclass so the type variable is overly generic.

import enum
import itertools
from typing import TypeVar

_EnumMemberT = TypeVar("_EnumMemberT")


class _IndexableEnumType(enum.EnumType):
    def __getitem__(self: type[_EnumMemberT], name: str | int) -> _EnumMemberT:
        if isinstance(name, int):
            try:
                return next(itertools.islice(enum.EnumType.__iter__(self), name, None))
            except StopIteration:
                raise IndexError("enum index out of range") from None
        return enum.EnumType.__getitem__(self, name)

    def index(self: type[_EnumMemberT], name: str) -> int:
        for index, member in enumerate(enum.EnumType.__iter__(self)):
            if member.name == name:  # type: ignore[attr-defined]
                return index
        raise ValueError(f"'{name}' is not in enum")


class IndexableEnum(enum.Enum, metaclass=_IndexableEnumType):
    pass

Then one can declare an IndexableEnum as per usual and use indices as you would with a list:

>>> class MyEnum(IndexableEnum):
...     Foo = "bar"
...     Baz = "qux"
... 
>>> 
>>> MyEnum[0] 
<MyEnum.Foo: 'bar'>
>>>                    
>>> MyEnum.index('Foo')
0

Answer 2

class Level(Enum):

    prerequisite_level: Optional["Level"]
    dependent_level: Optional["Level"]
    lower_priority_levels: List["Level"]
    greater_priority_levels: List["Level"]

    DATA_CHECK = "data check"
    DESIGN_CHECK = "design check"
    ALERT = "alert"

I'm having a hard time understanding the above: ... [comments clarified that the first four should be attributes, and prequisite and dependent are the previous and following members, respectively].

The solution is to modify previous members as the current member is being initialized (the trick being that the current member isn't added to the parent Enum until after the member's creation and initialization). Here is the solution using the stdlib's Enum¹ (Python 3.6 and later):

from enum import Enum, auto

class Level(str, Enum):
    #
    def __init__(self, name):
        # create priority level lists
        self.lower_priority_levels = list(self.__class__._member_map_.values())
        self.greater_priority_levels = []
        # update previous members' greater priority list
        for member in self.lower_priority_levels:
            member.greater_priority_levels.append(self)
        # and link prereq and dependent
        self.prerequisite = None
        self.dependent = None
        if self.lower_priority_levels:
            self.prerequisite = self.lower_priority_levels[-1]
            self.prerequisite.dependent = self
    #
    def _generate_next_value_(name, start, count, last_values, *args, **kwds):
        return (name.lower().replace('_',' '), ) + args
    #
    DATA_CHECK = auto()
    DESIGN_CHECK = auto()
    ALERT = auto()

and in use:

>>> list(Level)
[<Level.DATA_CHECK: 'data check'>, <Level.DESIGN_CHECK: 'design check'>, <Level.ALERT: 'alert'>]

>>> Level.DATA_CHECK.prerequisite
None

>>> Level.DATA_CHECK.dependent
<Level.DESIGN_CHECK: 'design check'>

>>> Level.DESIGN_CHECK.prerequisite
<Level.DATA_CHECK: 'data check'>

>>> Level.DESIGN_CHECK.dependent
<Level.ALERT: 'alert'>

>>> Level.ALERT.prerequisite
<Level.DESIGN_CHECK: 'design check'>

>>> Level.ALERT.dependent
None

Note: If you don't want to see the name twice, a custom __repr__ can show just the enum and member names:

    def __repr__(self):
        return '<%s.%s>' % (self.__class__.__name__, self.name)

then you'll see:

>>> Level.DESIGN_CHECK
<Level.DESIGN_CHECK>

---

¹If using Python 3.5 or older you need to use aenum².

² Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

Answer 3

You can subclass EnumMeta to override the __getitem__ method with additional conditions to return a list of Enum values or a specific Enum value based on the given index, and create a subclass of Enum with the aforementioned subclass of EnumMeta as the metaclass, so that any subclass of this new subclass of Enum can be indexed as desired:

from itertools import islice
from enum import Enum, EnumMeta

class IndexableEnumMeta(EnumMeta):
    def __getitem__(cls, index):
        if isinstance(index, slice):
            return [cls._member_map_[i] for i in islice(cls._member_map_, index.start, index.stop, index.step)]
        if isinstance(index, int):
            return cls._member_map_[next(islice(cls._member_map_, index, index + 1))]
        return cls._member_map_[index]

class IndexableEnum(Enum, metaclass=IndexableEnumMeta):
    pass

class Level(IndexableEnum):
    DATA_CHECK = "data check"
    DESIGN_CHECK = "design check"
    ALERT = "alert"

so that Level[1:3] returns:

[<Level.DESIGN_CHECK: 'design check'>, <Level.ALERT: 'alert'>]

and Level[1] returns:

Level.DESIGN_CHECK

(Credit goes to @EthanFurman for pointing out the viability of subclassing EnumMeta.)

Answer 4

A possible alternative to achieve the same result in terms of usage would be to use collections.namedtuple instead:

from collections import namedtuple
LevelSequence = namedtuple('Level', ('DATA_CHECK', 'DESIGN_CHECK', 'ALERT'))
Level = LevelSequence('data check', 'design check', 'alert')

So that:

• Level.DESIGN_CHECK and Level[1] both return 'design check', and
• Level[1:3] returns ('design check', 'alert')