如何为IntEnum类型的返回值添加类型提示？

问题场景 #

你希望给函数添加类型注解，表明返回值是某个IntEnum的成员，但直接返回原始整数（或从bytes中提取单个元素）时，Mypy会抛出类型不兼容的错误。比如以下示例：

from enum import IntEnum
import random

class GoodIntegers(IntEnum):
    THE_ANSWER_TO_LIFE_THE_UNIVERSE_AND_EVERYTHING = 42
    DEADBEEF = 0xdeadbeef
    LEET = 1337
    
def f() -> GoodIntegers:
    return random.choice([42, 1337, 0xdeadbeef])

Mypy报错：

test.py:12: error: Incompatible return value type (got "int", expected "GoodIntegers")  [return-value]
Found 1 error in 1 file (checked 1 source file)

实际场景中，你会从bytes字符串提取元素作为枚举值，类似这样的代码也会触发相同错误：

from enum import IntEnum
from typing import List, Tuple, Optional

class OpCodes(IntEnum):
    ADD = 1
    SUB = 2
    # 其他操作码定义...

def get_ops_from_script(script: bytes) -> List[Tuple[OpCodes, Optional[bytes]]]:
    ops = []
    n = 0
    while n < len(script):
        op: OpCodes = script[n]  # Mypy会报错：int无法赋值给OpCodes
        # 提取操作码对应的数据逻辑...
        n += 1
        ops.append((op, None))
    return ops

解决方案 #

1. 直接返回/使用IntEnum成员 #

如果你的整数来源可以替换为枚举成员，这是最安全的方式，Mypy能直接识别类型：

def f() -> GoodIntegers:
    return random.choice([
        GoodIntegers.THE_ANSWER_TO_LIFE_THE_UNIVERSE_AND_EVERYTHING,
        GoodIntegers.DEADBEEF,
        GoodIntegers.LEET
    ])

2. 使用类型断言（cast） #

当你确定整数一定属于枚举的有效值范围，但无法直接使用枚举成员时，可以用typing.cast告诉Mypy强制转换类型：

from typing import cast

def get_ops_from_script(script: bytes) -> List[Tuple[OpCodes, Optional[bytes]]]:
    ops = []
    n = 0
    while n < len(script):
        op = cast(OpCodes, script[n])
        # 后续逻辑...
        n += 1
        ops.append((op, None))
    return ops

注意：这种方式没有运行时检查，如果实际值不在枚举定义中，运行时不会报错，可能引发后续逻辑问题。

3. 用IntEnum构造函数做运行时验证 #

如果需要同时满足类型检查和运行时合法性验证，直接用枚举类构造函数转换整数：

def get_ops_from_script(script: bytes) -> List[Tuple[OpCodes, Optional[bytes]]]:
    ops = []
    n = 0
    while n < len(script):
        op = OpCodes(script[n])  # 运行时若值无效会抛出ValueError
        # 后续逻辑...
        n += 1
        ops.append((op, None))
    return ops

这种方式Mypy会自动识别op的类型为OpCodes，同时保证运行时只有合法的枚举值会被使用，安全性最高，但会带来少量构造枚举的性能开销。

内容的提问来源于stack exchange，提问作者PiRK