Python - получить значение для enum по умолчанию при использовании в сравнении

Question

Я создал класс Enum как показано:

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"

Я хочу использовать этот класс для сравнения со строкой, которую я получаю после прочтения сообщения. Я сравниваю значения, как показано. Значение в msg_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)

Я ожидаю, что для msg_type = "A" оператор msg_type == ct.MsgType.LOGON должен быть True, но вместо этого выполняется оператор else.

Если я напишу ct.MsgType.LOGON.value, то получу желаемый результат. Но я хочу, чтобы это поведение было по умолчанию для класса. Какой метод я должен переопределить или я должен попробовать другой подход?

Answer 1

Ключевой вопрос здесь - что должно произойти, если msg_type недействительноНапример, либо из-за ошибки, либо из-за повреждения msg_type == 'z'.

Если исключение подходит, то проще всего его получить с помощью

msg_type = MsgType(msg_type)

, который сгенерирует ValueErrorдля недопустимых значений, и вы можете либо немедленно разобраться с этим:

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

, либо позволить ему перейти к более высокой функции.

Если вы не хотите / нуждаетесь в исключениях, тогда другойОбходное сравнение с MsgType.LOGON.value означает, что MsgType Enum также является производным от str:

class MsgType(str, Enum):
    # definitions here

Тогда в вашем коде вы можете сказать:

if msg_type == 'A':
    # stuff

без преобразования msg_type в MsgType.

Я предпочитаю первый метод (не смешивать в str и не выполнять преобразование), если вы не добавляете Enum к существующемукодовая база и LOGON, HEARTBEAT и т. д. являются константами и уже используются в других местах.

Обратите внимание, что даже со вторым методом вы все равно можете использовать MsgType(msg_type), чтобы вызвать исключение ValueError, если вынадо.

Answer 2

msg_type = "A" - это значение в Enum.Вам нужно изменить одну сторону ваших равных сравнений.Либо:

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

, либо это:

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

РЕДАКТИРОВАТЬ: я вижу ваше заявление об использовании .value сейчас.Вот как работают эти перечисления ... если вы хотите переопределить __eq__ в вашем MsgType классе, вы можете.Но в процессе вы нарушите сравнение равенства по умолчанию для Enums, и вам потребуется специальная проверка типов, чтобы проверить, является ли ваша левая / правая сторона строкой / Enum / и т.д.Я бы просто заставил вашу read_header функцию возвращать экземпляр Enum, а не строковое значение.