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

Issue

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?

Solution

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.

Answered By - Ethan Furman