Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Implement comparison operator for Enum type in python3

I'm really beginner in Python, so please forgive me if this is something obvious.

I have an enum class, and I want to be able to compare the members. Below code seems to do what I want ( but not how I want )

import enum

class AB(enum.Enum):

    a=1
    b=2
    c=3
    d=4
    e=5

    @classmethod
    def le(cls, a, b):
        lst = [cls.a, cls.b, cls.c, cls.d, cls.e]
        ia = lst.index(a)
        ib = lst.index(b)
        return(ia <= ib)


if AB.le(AB.a, AB.b):
    print('Do this')
else:
    print('Do that')

Now my question is how to code the comparison operator __le__ so that I can run the below code instead?

mem1 = AB.a
mem2 = AB.c

if mem1 <= mem2 :
    print('Do this')
else:
    print('Do that')
like image 551
kollery Avatar asked Dec 24 '22 02:12

kollery


1 Answers

Enum subclasses are somewhat special in that all enumeration values become instances of the class (with a few tweaks). This does mean you can 'just' define a normal method on the Enum subclass and they'll be available on each enumeration value.

This applies to special methods like object.__le__() too; just define it as a regular method, not a classmethod:

class AB(enum.Enum):
    def __le__(self, b):
        return self.value <= b.value

    a = 1
    b = 2
    c = 3
    d = 4
    e = 5

Note that I used the instance attribute .value, just like you can do AB.a.value.

You could also use the IntEnum class; that makes each enumeration value a subclass of int, and they can be compared naturally:

class AB(enum.IntEnum):
    a = 1
    b = 2
    c = 3
    d = 4
    e = 5

Demo:

>>> import enum
>>> class AB(enum.Enum):
...     def __le__(self, b):
...         return self.value <= b.value
...     a = 1
...     b = 2
...     c = 3
...     d = 4
...     e = 5
...
>>> AB.a <= AB.b
True
>>> AB.b <= AB.a
False
>>> AB.a < AB.b
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unorderable types: AB() < AB()

and using IntEnum as the base:

>>> class AB(enum.IntEnum):
...     a = 1
...     b = 2
...     c = 3
...     d = 4
...     e = 5
...
>>> AB.a <= AB.b
True
>>> AB.b >= AB.a
True
>>> AB.b > AB.a
True
>>> AB.a + AB.b
3
like image 172
Martijn Pieters Avatar answered Jan 04 '23 23:01

Martijn Pieters