How to get the Enum value by default when used in a comparison?

Question

I have created a Enum class as shown:

class MsgType(Enum):
    # ADMINISTRATIVE MESSAGE
    HEARTBEAT = "0"
    LOGON = "A"
    LOGOUT = "5"
    REJECT_SESSION_LEVEL = "3"
    RESEND_REQUEST = "2"
    SEQUENCE_RESET = "4"
    SESSION_REJECT = "3"
    TEST_REQUEST = "1"

I want to use this class for comparison with a string that I get after reading a message. I am comparing the values as shown. The value in msg_type is a of type str.

def read_admin_msg(message):
    msg_type = read_header(message)
    if msg_type == ct.MsgType.HEARTBEAT:
        print(msg_type)
    elif msg_type == ct.MsgType.LOGON:
        print(msg_type)
    elif msg_type == ct.MsgType.LOGOUT:
        print(msg_type)
    elif msg_type == ct.MsgType.REJECT_SESSION_LEVEL:
        print(msg_type)
    elif msg_type == ct.MsgType.RESEND_REQUEST:
        print(msg_type)
    elif msg_type == ct.MsgType.SEQUENCE_RESET:
        print(msg_type)
    elif msg_type == ct.MsgType.SESSION_REJECT:
        print(msg_type)
    elif msg_type == ct.MsgType.TEST_REQUEST:
        print(msg_type)
    else:
        print("Not found")
        print(msg_type)

My expectation is that for msg_type = "A" the statement msg_type == ct.MsgType.LOGON should be True but instead else statement is executed.

If I write ct.MsgType.LOGON.value then I get the desired result. But I want this behavior to be default for the class. Which method should I override or should I try a different approach?

It returns a `str` and it's value is one of the values in `MsgType()` class — ap14, Sep 14 '18 at 11:45
I actually would like to return enum value when i call it like this `ct.MsgType.LOGON` but there is no solution to that — ruslaniv, Nov 03 '22 at 17:23

wholevinski · Accepted Answer · 2018-09-14T11:42:59.317

6

msg_type = "A" is the value of the Enum. You need to change one side of your equal comparisons. Either:

elif msg_type == MsgType.LOGON.value:
    print(msg_type)

or this:

# This one is probably preferred
elif MsgType(msg_type) == MsgType.LOGON:
    print(msg_type)

EDIT: I see your statement about using .value now. That's just how these enums work...if you want to override the __eq__ on your MsgType class, you can. But you'd be breaking the default equality comparison for Enums in the process, and you'd have to have special type checking in there for checking if your left/right side is a string/Enum/etc. I'd just make your read_header function return an instance of the Enum, and not the string value.

edited Sep 14 '18 at 11:42

answered Sep 14 '18 at 11:37

wholevinski

3,658
17
23

2

In the second case, probably better to do it when getting the value: `msg_type = ct.MsgType(read_header(message))`. – Daniel Roseman Sep 14 '18 at 11:42
@DanielRoseman Heh, yep, I was just putting that into my edit I just posted. – wholevinski Sep 14 '18 at 11:43

score 3 · Answer 2 · edited Aug 25 '22 at 11:10

The key question here is what should happen if msg_type is invalid. For example, either because of a bug or corruption msg_type == 'z'.

If an exception is appropriate then the easiest way to get it is with

msg_type = MsgType(msg_type)

which will generate a ValueError for illegal values, and you can either deal with it immediately:

try:
    msg_type = MsgType(msg_type)
    # your if...elif...else here
except ValueError:
    # handle problem here

or let it bubble up to a higher function.

If you don't want/need exceptions, then another way around around having to compare with MsgType.LOGON.value is to make the MsgType Enum derive from str as well:

class MsgType(str, Enum):
    # definitions here

Then in your code you can say:

if msg_type == 'A':
    # stuff

without converting msg_type to MsgType.

My preference is the first method (not mixing in str, and doing the conversion) unless you are adding Enum to an existing code base and LOGON, HEARTBEAT, etc. are constants and already in use elsewhere.

Note that even with the second method you can still use MsgType(msg_type) to raise a ValueError exception if you need to.

I do like this solution, `class MsgType(str, Enum):` – westonplatter Sep 16 '22 at 21:00 — westonplatter, Sep 16 '22 at 21:00

Christopher Peisert · Answer 3 · 2022-05-10T17:19:30.337

Override `eq` to check either `Enum` or `Enum.value`

For Enums with string values, one solution to avoid explicitly calling Enum.value is to override __eq__ to check if the Enum is being compared to a string, and if not, then use the default Enum equality check.

def __eq__(self, other):
    if isinstance(other, str):
        return self.value == other

    if isinstance(other, Enum):
        return self.value == other.value

    return False

Note that this works if the Enum is on the left-hand or right-hand side of the equality.

Example

>>> from enum import Enum
>>> class MyEnum(Enum):
        BAR = 'bar'
        def __eq__(self, other):
            if isinstance(other, str):
                return self.value == other
            if isinstance(other, Enum):
                return self.value == other.value
            return False
   
>>> MyEnum.BAR == MyEnum.BAR
True
>>> MyEnum.BAR == 'bar'
True
>>> 'bar' == MyEnum.BAR
True
>>> MyEnum.BAR.value == 'bar'
True
>>> MyEnum.BAR == 'foo'
False

How to get the Enum value by default when used in a comparison?

3 Answers3

Override __eq__ to check either Enum or Enum.value

Example

Override `eq` to check either `Enum` or `Enum.value`