How can I elegantly find the next and previous value in a Python Enum? [duplicate]

Question

I have a simple enum in Python that looks like this:

from enum import Enum

class MyEnum(Enum):
    #All members have increasing non-consecutive integer values.
    A = 0
    B = 2
    C = 10
    D = 18
    ...

I want functions pred() and succ() that given a member of MyEnum return the member of MyEnum that precedes and succeeds the given element, respectively (just like the functions of the same name in Haskell). For example succ(MyEnum.B) and pred(MyEnum.D) should both return MyEnum.C. An exception can be raised if succ is called on the last member of pred is called on the first member.

There doesn't seem to be any built in method to do this, and while I can call iter(MyEnum) to iterate over the values it has to go through the whole enum from the beginning. I could probably implement a sloppy loop to accomplish this on my own, but I know there are some real Python gurus on this site, so to you I ask: is there a better approach?

Answer 1

Note that you can provide succ and pred methods inside an Enum class:

class Sequential(Enum):
    A = 1
    B = 2
    C = 4
    D = 8
    E = 16

    def succ(self):
        v = self.value * 2
        if v > 16:
            raise ValueError('Enumeration ended')
        return Sequential(v)

    def pred(self):
        v = self.value // 2
        if v == 0:
            raise ValueError('Enumeration ended')
        return Sequential(v)

Used as:

>>> import myenums
>>> myenums.Sequential.B.succ()
<Sequential.C: 4>
>>> myenums.Sequential.B.succ().succ()
<Sequential.D: 8>
>>> myenums.Sequential.B.succ().succ().pred()
<Sequential.C: 4>

Obviously this is efficient only if you actually have a simple way to compute the values from an item to the next or preceding one, which may not always be the case.

If you want a general efficient solution at the cost of adding some space, you can build the mappings of the successor and predecessor functions. You have to add these as attributes after the creation of the class (since Enum messes up attributes) so you can use a decorator to do that:

def add_succ_and_pred_maps(cls):
    succ_map = {}
    pred_map = {}
    cur = None
    nxt = None
    for val in cls.__members__.values():
        if cur is None:
            cur = val
        elif nxt is None:
            nxt = val

        if cur is not None and nxt is not None:
            succ_map[cur] = nxt
            pred_map[nxt] = cur
            cur = nxt
            nxt = None
    cls._succ_map = succ_map
    cls._pred_map = pred_map

    def succ(self):
        return self._succ_map[self]

    def pred(self):
        return self._pred_map[self]

    cls.succ = succ
    cls.pred = pred
    return cls





@add_succ_and_pred_maps
class MyEnum(Enum):
    A = 0
    B = 2
    C = 8
    D = 18

Used as:

>>> myenums.MyEnum.A.succ()
<MyEnum.B: 2>
>>> myenums.MyEnum.B.succ()
<MyEnum.C: 8>
>>> myenums.MyEnum.B.succ().pred()
<MyEnum.B: 2>
>>> myenums.MyEnum._succ_map
{<MyEnum.A: 0>: <MyEnum.B: 2>, <MyEnum.C: 8>: <MyEnum.D: 18>, <MyEnum.B: 2>: <MyEnum.C: 8>}

you probably want a custom exception instead of KeyError but you get the idea.

---

There probably is a way to integrate the last step using metaclasses, but it's notstraightforward for the simple fact thatEnums are implemented using metaclasses and it's not trivial to compose metaclasses.