如何让Pyright通过枚举判断自动推断类型？

嘿，这个问题我之前也折腾过，完全可以不用cast()或者isinstance()就让Pyright精准推断出类型！核心思路是给每个数据类的type字段绑定具体的枚举成员类型，而不是笼统的枚举类本身，这样Pyright就能通过枚举值的相等判断自动缩小类型范围。

直接看修改后的代码：

from dataclasses import dataclass
from enum import Enum, auto
from typing import Literal

class Type(Enum):
    FOO = auto()
    BAR = auto()

@dataclass
class Foo:
    # 把type字段的类型指定为具体的枚举成员Type.FOO
    type: Literal[Type.FOO]

@dataclass
class Bar:
    # 同理，绑定Type.BAR
    type: Literal[Type.BAR]

item = next(i for i in (Foo(Type.FOO), Bar(Type.BAR)) if i.type == Type.BAR)
reveal_type(item)  # 现在Pyright会正确推断为Bar，不再是Foo | Bar！

原理说明 #

原来的代码里，Foo和Bar的type字段都是Type类型，Pyright没法区分哪个类对应哪个枚举值。当我们用Literal[Type.FOO]和Literal[Type.BAR]分别注解后，Pyright就明确知道：

• 只有Foo实例的type会是Type.FOO
• 只有Bar实例的type会是Type.BAR

所以当你用i.type == Type.BAR过滤时，它能直接推断出符合条件的只能是Bar类型的实例。

另外还有一种适配复杂场景的方法——用类型守卫函数，同样不需要isinstance()：

from typing import TypeGuard

# 定义一个类型守卫，告诉Pyright当返回True时，item是Bar类型
def is_bar(item: Foo | Bar) -> TypeGuard[Bar]:
    return item.type == Type.BAR

# 用类型守卫过滤
item = next(i for i in (Foo(Type.FOO), Bar(Type.BAR)) if is_bar(i))
reveal_type(item)  # 同样会推断为Bar

这种方法适合枚举和类型的对应关系更复杂的场景，不过第一种直接绑定Literal类型的方式更简洁直观。

备注：内容来源于stack exchange，提问作者Q42