How to statically protect against str == enum comparisons?

Question

More often than not, I make the following mistake:

import enum

class StatusValues(enum.Enum):
    one = "one"
    two = "two"

def status_is_one(status: str):
    return status == StatusValues.one

String will never be an enum class. The problem is that it should be StatusValues.one.value.

Is there a strcmp(status, StatusValues.one)-ish function so that my pyright will error on the line that I am comparing string with a class? Is there a good way to protect against such mistakes?

Answer 1

Yes, this can be a gotcha: by default we cannot directly compare strings with enum members:

>>> status = "two"
>>> status == StatusValues.two
False

One has to remember to compare a string to the enum member's .value, which is also kind of verbose:

>>> "two" == StatusValues.two.value
True

Fortunately, there is a solution, which is even mentioned in the docs, but it used to be (pre- Python 3.11) in a section that was somewhat easy to miss (IMHO). We can mix in a type of enum members:

class IntEnum(int, Enum):
    pass

This demonstrates how similar derived enumerations can be defined; for example a StrEnum that mixes in str instead of int.

The solution is therefore to define an enum with a str base class:

class StatusValues(str, Enum):
    one = "one"
    two = "two"

>>> status = "two"
>>> status == StatusValues.two  # no .value
True

In Python 3.11+ StrEnum is already included in the enum module, but you mentioned that you still need to support Python 3.7.