pythonenumscomparisonmembership

Check if variable of type Enum is in a set of Enums


According to the documentation of Enums in Python, comparison primarily refer to the is operator, due to Enums being a singleton. However, if one would check if an Enum is in a set of Enums, is it then perfectly fine to use the in operator? It seems to work as expected, but my knowledge comes short here if there are any caveats that must be taken.

The following code should suffice as an illustrative example:

from enum import auto, StrEnum


class MyEnum(StrEnum):
    STR1 = auto()
    STR2 = auto()
    STRN = auto()


def main():

    foo = MyEnum.STR1

    if foo in {MyEnum.STR1, MyEnum.STR2}: # is this okay?
        print("foo")

    if foo is MyEnum.STR1 or foo is MyEnum.STR2: # the alternative following the documentation
        print("foo")

main()

I have looked up the documentation, and found the is operator to be best practice, as well as related questions here on StackOverflow. Moreover, I have experimented with examples my self, which seems to work as expected. Nevertheless, I did not succeed finding an answer to my exact question.

Any insight to whether if it is okay, and also, what the actual difference might be at a lower level is greatly appreciated.


Solution

  • Generally, if a is b, then a == b. This is the reflexive property of equality: an object should equal to itself.

    Generally, if b is an iterable and a in b, then any(value == a for value in b). If a is contained in b, then b must have at least one element that is equal to a.

    Therefore, you can assume

    foo is MyEnum.STR1 or foo is MyEnum.STR2
    

    is equivalent to

    foo == MyEnum.STR1 or foo == MyEnum.STR2
    

    which is equivalent to

    foo in {MyEnum.STR1, MyEnum.STR2}
    

    Using them interchangeably is safe, since both set and enums are built-in types and do not have any special logic for equality or membership.

    Of course checking using is is a good practice, but readability of your code is a better practice. If you need to compare with five enum values, which code is more readable?

    if foo in {MyEnum.STR1, MyEnum.STR2, MyEnum.STR3, MyEnum.STR4, MyEnum.STR5}:
        pass
    
    if foo is MyEnum.STR1 or foo is MyEnum.STR2 or foo is MyEnum.STR3 or foo is MyEnum.STR4 or foo is MyEnum.STR5:
        pass