How to compare Enums in Python?

Question

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!

Solution 1

I hadn'r encountered Enum before so I scanned the doc (https://docs.python.org/3/library/enum.html) ... and found OrderedEnum (section 8.13.13.2) Isn't this what you want? From the doc:

>>> class Grade(OrderedEnum):
...     A = 5
...     B = 4
...     C = 3
...     D = 2
...     F = 1
...
>>> Grade.C < Grade.A
True

Solution 2

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
>>> 

Solution 3

Combining some of the above ideas, you can subclass enum.Enum to make it comparable to string/numbers and then build your enums on this class instead:

import numbers
import enum


class EnumComparable(enum.Enum):
    def __gt__(self, other):
        try:
            return self.value > other.value
        except:
            pass
        try:
            if isinstance(other, numbers.Real):
                return self.value > other
        except:
            pass
        return NotImplemented

    def __lt__(self, other):
        try:
            return self.value < other.value
        except:
            pass
        try:
            if isinstance(other, numbers.Real):
                return self.value < other
        except:
            pass
        return NotImplemented

    def __ge__(self, other):
        try:
            return self.value >= other.value
        except:
            pass
        try:
            if isinstance(other, numbers.Real):
                return self.value >= other
            if isinstance(other, str):
                return self.name == other
        except:
            pass
        return NotImplemented

    def __le__(self, other):
        try:
            return self.value <= other.value
        except:
            pass
        try:
            if isinstance(other, numbers.Real):
                return self.value <= other
            if isinstance(other, str):
                return self.name == other
        except:
            pass
        return NotImplemented

    def __eq__(self, other):
        if self.__class__ is other.__class__:
            return self == other
        try:
            return self.value == other.value
        except:
            pass
        try:
            if isinstance(other, numbers.Real):
                return self.value == other
            if isinstance(other, str):
                return self.name == other
        except:
            pass
        return NotImplemented

Solution 4

You can create a simple decorator to resolve this too:

from enum import Enum
from functools import total_ordering

def enum_ordering(cls):
    def __lt__(self, other):
        if type(other) == type(self):
            return self.value < other.value

        raise ValueError("Cannot compare different Enums")

    setattr(cls, '__lt__', __lt__)
    return total_ordering(cls)


@enum_ordering
class Foos(Enum):
    a = 1
    b = 3
    c = 2

assert Names.a < Names.c
assert Names.c < Names.b
assert Names.a != Foos.a
assert Names.a < Foos.c # Will raise a ValueError

For bonus points you could implement the other methods in @VoteCoffee's answer above

View more solutions

Comments

[Sebastian Werk]

Great description, thanks a lot. Just one question: You return NotImplemented instead of raise NotImplemented. Is there a general rule, when to use return and when raise?

[juanpa.arrivillaga]

@SebastianWerk Well, you cannot raise NotImplemented because it is not an exception. It is a built-in singleton. See the docs, it is there for the special case of the rich-comparison operators. The NotImplementedError, according to the docs, is there for when "abstract methods should raise this exception when they require derived classes to override the method. ".

[juanpa.arrivillaga]

@SebastianWerk Also, see this question: stackoverflow.com/questions/878943/…

[Cecil Curry]

Excellent answer, good Sir. This approach is a succinct – albeit less efficient – alternative to the OrderedEnum class detailed in the official Python documentation. While the OrderedEnum solution of manually implementing all comparison operators is modestly faster, the @total_ordering solution given above has its merits. Brevity is a thankless virtue. Relatedly, does anyone know why OrderedEnum was merely documented rather than added to the enum module?

[asu]

This approach doesn't work due to unfixed Python bug: bugs.python.org/issue30545 You should always override comparator methods like eq and lt

[juanpa.arrivillaga]

@asu there is no unfixed bug, the issues in that thread are not a problem with Enum itself, and in any case, this example does override the comparison operators.

[Cobertos]

Can't seem to import in Python 3.6, ImportError: cannot import name 'OrderedEnum'? EDIT: It looks like this is an "Interesting Example" and not actually in the standard python library. You'd need to copy the snippet from the docs to use it.

[WestCoastProjects]

Is there a python Enum construct or similar that actually kinda .. works ? Adding @total_ordering and implementing comparison operators is not something should have to code ourselves. This is too much boilerplate.