如何根据函数参数实际值定义返回类型？Pyright类型推断疑问

问题描述 #

已知可通过TypeVar根据参数类型定义函数返回类型，但希望根据参数内容来定义。以open函数为例：

• 传入mode为"r"时，Pyright推断返回TextIOWrapper
• 传入"rb"时，推断为BufferedReader
• 当mode是字面量赋值的变量时也能正确推断
• 但当mode是类型为str的函数参数时，Pyright推断为IO[Any]

疑问：Pyright的该推断逻辑是基于open的类型提示，还是对open做了硬编码？

解答 #

Pyright对open函数的返回类型推断，核心基于标准库中open函数的类型注解，而非硬编码，但这个注解用到了Python类型系统的两个关键特性：

1. 重载（@overload）与字面量类型 #

标准库的open函数通过@overload装饰器定义了多组精确的函数签名，每组签名对应特定的mode字面量值，并明确指定对应的返回类型。示例简化版如下：

from typing import overload, Literal, TextIOWrapper, BufferedReader, IO

@overload
def open(file: str, mode: Literal["r", "rt", "r+"], encoding: str | None = ...) -> TextIOWrapper: ...

@overload
def open(file: str, mode: Literal["rb", "rb+"], encoding: None = ...) -> BufferedReader: ...

# 兜底重载，匹配所有其他情况
@overload
def open(file: str, mode: str, encoding: str | None = ...) -> IO[Any]: ...

2. 类型推断的匹配逻辑 #

• 当传入字面量字符串（如"r"）或被标注为特定字面量类型的变量（如mode: Literal["rb"] = "rb"）时，Pyright能精准匹配对应的重载签名，从而推断出具体的返回类型。
• 当mode是未限定字面量范围的普通str类型参数时，Pyright无法确定它会匹配哪一组重载，因此会 fallback 到所有重载返回类型的公共父类——IO[Any]。

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