如何让mypy识别FlatBuffer生成的Python枚举类而非int类型？

问题描述 #

FlatBuffer生成的Python枚举类是普通类，成员为int类型：

class ErrorType(object):
    ERROR1 = 0
    ERROR2 = 1

当将其用作函数参数类型注解时：

def my_function(error: ErrorType):
    pass

调用my_function(ErrorType.ERROR1)时代码可正常运行，但mypy会报错，判定传入的是int类型而非ErrorType类型：

test.py:8: error: Argument 1 to "my_function" has incompatible type "int"; expected "ErrorType"  [arg-type]
Found 1 error in 1 file (checked 1 source file)

解决方法 #

方法一：修改FlatBuffer生成配置，生成原生Python枚举 #

FlatBuffer支持生成继承自enum.Enum的原生枚举类，生成代码时添加对应参数即可：

flatc --python --python-enum-type=enum.Enum your_schema.fbs

生成后的代码会自动变为：

from enum import Enum

class ErrorType(Enum):
    ERROR1 = 0
    ERROR2 = 1

此时ErrorType.ERROR1的类型就是ErrorType，mypy可正确识别，函数调用不会报错。

方法二：手动修改生成的枚举类 #

如果无法重新生成代码，可手动修改ErrorType类让它继承自enum.Enum：

from enum import Enum

class ErrorType(Enum):
    ERROR1 = 0
    ERROR2 = 1

修改后，枚举成员的类型会从int变为ErrorType实例，符合类型注解要求。

方法三：调整函数的类型注解（无需修改枚举类） #

如果不能修改枚举类，可用typing.Literal限定参数只能是枚举的合法值：

from typing import Literal

def my_function(error: Literal[ErrorType.ERROR1, ErrorType.ERROR2]):
    pass

这种方式下，mypy会允许传入ErrorType.ERROR1或ErrorType.ERROR2，同时阻止传入其他无效int值，达到类型检查目的。

也可以用Union，但检查精度不如Literal：

from typing import Union

def my_function(error: Union[int, ErrorType]):
    pass

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