How to compare Enums in Python?

Sebastian Werk picture Sebastian Werk · Sep 1, 2016 · Viewed 31.8k times · Source

Since Python 3.4, the Enum class exists.

I am writing a program, where some constants have a specific order and I wonder which way is the most pythonic to compare them:

class Information(Enum):
    ValueOnly = 0
    FirstDerivative = 1
    SecondDerivative = 2

Now there is a method, which needs to compare a given information of Information with the different enums:

information = Information.FirstDerivative
print(value)
if information >= Information.FirstDerivative:
    print(jacobian)
if information >= Information.SecondDerivative:
    print(hessian)

The direct comparison does not work with Enums, so there are three approaches and I wonder which one is preferred:

Approach 1: Use values:

if information.value >= Information.FirstDerivative.value:
     ...

Approach 2: Use IntEnum:

class Information(IntEnum):
    ...

Approach 3: Not using Enums at all:

class Information:
    ValueOnly = 0
    FirstDerivative = 1
    SecondDerivative = 2

Each approach works, Approach 1 is a bit more verbose, while Approach 2 uses the not recommended IntEnum-class, while and Approach 3 seems to be the way one did this before Enum was added.

I tend to use Approach 1, but I am not sure.

Thanks for any advise!

Answer

juanpa.arrivillaga picture juanpa.arrivillaga · Sep 1, 2016

You should always implement the rich comparison operaters if you want to use them with an Enum. Using the functools.total_ordering class decorator, you only need to implement an __eq__ method along with a single ordering, e.g. __lt__. Since enum.Enum already implements __eq__ this becomes even easier:

>>> import enum
>>> from functools import total_ordering
>>> @total_ordering
... class Grade(enum.Enum):
...   A = 5
...   B = 4
...   C = 3
...   D = 2
...   F = 1
...   def __lt__(self, other):
...     if self.__class__ is other.__class__:
...       return self.value < other.value
...     return NotImplemented
... 
>>> Grade.A >= Grade.B
True
>>> Grade.A >= 3
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unorderable types: Grade() >= int()

Terrible, horrible, ghastly things can happen with IntEnum. It was mostly included for backwards-compatibility sake, enums used to be implemented by subclassing int. From the docs:

For the vast majority of code, Enum is strongly recommended, since IntEnum breaks some semantic promises of an enumeration (by being comparable to integers, and thus by transitivity to other unrelated enumerations). It should be used only in special cases where there’s no other choice; for example, when integer constants are replaced with enumerations and backwards compatibility is required with code that still expects integers.

Here's an example of why you don't want to do this:

>>> class GradeNum(enum.IntEnum):
...   A = 5
...   B = 4
...   C = 3
...   D = 2
...   F = 1
... 
>>> class Suit(enum.IntEnum):
...   spade = 4
...   heart = 3
...   diamond = 2
...   club = 1
... 
>>> GradeNum.A >= GradeNum.B
True
>>> GradeNum.A >= 3
True
>>> GradeNum.B == Suit.spade
True
>>>