在 Python 类型中表达枚举与其成员之间的关系

Question

如何键入（在 Python 中，例如对于 MyPy）一个需要两个参数的函数——一个枚举和它的值/成员之一？

from enum import Enum
from typing import TypeVar, Type

class MyEnumA(Enum):
   A = 1
   B = 2

class MyEnumB(Enum):
   A = 1
   B = 2

TE = TypeVar('TE', bound=Enum)
def myfunction(member: TE, e: Type[TE]) -> None:
    pass

myfunction(MyEnumA.A, MyEnumA) # all right
myfunction(MyEnumA.A, MyEnumB) # I expect mypy-error here but it passed

print(type(MyEnumA.A)) # says: <enum 'MyEnumA'>
print(type(MyEnumB.A)) # says: <enum 'MyEnumB'>

print(f'{isinstance(MyEnumA.A, MyEnumA)=}') # says: isinstance(MyEnumA.A, MyEnumA)=True
print(f'{isinstance(MyEnumA.A, MyEnumB)=}') # says: isinstance(MyEnumA.A, MyEnumB)=False

reveal_type(MyEnumA) # mypy: Revealed type is "def (value: builtins.object) -> e.MyEnumA"
reveal_type(MyEnumA.A) # mypy: Revealed type is "Literal[e.MyEnumA.A]?"

我想了解

• 为什么 MyPy 没有为第二次调用声明错误；和
• 如何键入

  myfunction

  以便 MyPy 在那里检测到错误。

Answer 1

关键观察是 MyPy 实际上认为类 MyEnumA 是什么：

reveal_type(MyEnumA) # mypy: Revealed type is "def (value: builtins.object) -> e.MyEnumA"

它给出了注释

myfunction

的线索，无论它多么不明显：

P = ParamSpec('P')
def myfunction(member: TE, e: Callable[P, TE]) -> None:

这使得 MyPy 在

myfunction

的第二次调用中发现输入错误：

error: Argument 2 to "myfunction" has incompatible type "Type[MyEnumA]"; expected "Callable[[object], MyEnumB]"  [arg-type]

Answer 2

根据 mypy 文档，任何类都是有效类型，因此对于绑定类型变量，它应该只是

MyEnumA

这将绑定它只接受您的特定类型的 Enum 子类，而不是所有扩展基类的 Enum 类.然后引用你想要一个参数只接受不同类型的枚举，你会使用

Type[...]

，所以结果是这个代码：

from enum import Enum
from typing import TypeVar, Type

class MyEnumA(Enum):
   A = 1
   B = 2

class MyEnumB(Enum):
   A = 1
   B = 2

TE = TypeVar('TE', bound=MyEnumA)
def myfunction(member: TE, e: Type[TE]) -> None:
    pass

myfunction(MyEnumA.A, MyEnumA) # all right
myfunction(MyEnumB.A, MyEnumB) # I expect mypy-error here but it passed

您需要在哪里更改传递给

TypeVar

的枚举，但它会产生您想要的结果：

test.py:17: error: Value of type variable "TE" of "myfunction" cannot be "MyEnumB"  [type-var] Found 1 error in 1 file (checked 1 source file)