How can I check ordering of items in a Python Enum?

Viewed 942

I have an enum.Enum subclass:

class MyEnum(Enum):
    A = "apple"
    C = "cherry"
    B = "banana"

and I want to be able to use < and > to see if a given member comes before or after another one in the definition order, so I could do something like this:

>>> MyEnum.A < MyEnum.B
True
>>> MyEnum.B >= MyEnum.C
True
>>> MyEnum.C < MyEnum.A
False

based on where the values appear in the definitions of the enum, not the enum values themselves. I know that Enums preserve order, but there is no way of finding which came first. How can I accomplish this in Python 3.7?

5 Answers

You need to override comparison operators and somehow check the order of names of compared enum members. I found that _member_names_ preserve the order of defined members:

from enum import Enum
import functools


@functools.total_ordering
class MyEnum(Enum):
    A = "apple"
    C = "cherry"
    B = "banana"

    def __eq__(self, other):
        if isinstance(other, MyEnum):
            return (
                self._member_names_.index(self.name) ==
                self._member_names_.index(other.name)
            )
        return NotImplemented

    def __gt__(self, other):
        if isinstance(other, MyEnum):
            return (
                self._member_names_.index(self.name) >
                self._member_names_.index(other.name)
            )
        return NotImplemented


print(MyEnum.A < MyEnum.B)  # True
print(MyEnum.B >= MyEnum.C)  # True
print(MyEnum.C < MyEnum.A)  # False

You can use IntEnums instead and use better names:

import enum

class Fruits(enum.IntEnum):
    Apple = 0
    Cherry = 1
    Banana = 2

print(Fruits.Apple < Fruits.Banana) 
print(Fruits.Banana >= Fruits.Cherry) 
print(Fruits.Cherry < Fruits.Apple)

Output:

True
True
False

IntEnums can be compared to each other and normal ints - so Fruits.Apple < 1000 is valid.

You can get the name or value from it like this:

>>> print(Fruits.Apple.name)
Apple

>>> print(Fruits.Apple.value)
0

Using the OrderedEnum recipe as a start, it is easy to create your own ordered enum base class:

class OrderedEnum(Enum):

    def __init__(self, value, *args, **kwds):
        super().__init__(*args, **kwds)
        self.__order = len(self.__class__)

    def __ge__(self, other):
        if self.__class__ is other.__class__:
            return self.__order >= other.__order
        return NotImplemented

    def __gt__(self, other):
        if self.__class__ is other.__class__:
            return self.__order > other.__order
        return NotImplemented

    def __le__(self, other):
        if self.__class__ is other.__class__:
            return self.__order <= other.__order
        return NotImplemented

    def __lt__(self, other):
        if self.__class__ is other.__class__:
            return self.__order < other.__order
        return NotImplemented

And your example class becomes:

class MyEnum(OrderedEnum):
    A = "apple"
    C = "cherry"
    B = "banana"

you can override those operators in your enum class, it will be better for you to read about the data model here

if i would need to do it i would probably create a list of the elements ordered, and check it when calling does function, or change any element to be a tuple with index, and change the __get__ function

Assuming that your Enum class's __dict__ preserves Python 3.7's dictionary ordering semantics, you can get the order of the members like this:

>>> class E(enum.Enum):
...    A = 'apple'
...    C = 'cherry'
...    B = 'banana'
... 
>>> 
>>> members = [v for v in E.__dict__.values() if isinstance(v, E)]
>>> members
[<E.A: 'apple'>, <E.C: 'cherry'>, <E.B: 'banana'>]

The comparisons can then be made based on positions in the list.

Related