Python函数如何根据参数值而非类型动态改变返回值类型提示？

Python函数：根据参数值动态调整返回值类型提示

本文介绍如何在Python中创建函数，使其返回值类型提示根据参数值而非类型进行动态调整，从而提升IDE代码补全和类型检查的准确性。 这类似于内置open()函数：mode='r'返回str，mode='rb'返回bytes。 我们将学习如何实现这种特性。

传统类型提示依赖参数类型推断返回值类型。然而，我们的目标是基于参数值确定返回值类型。这需要用到Python的overload装饰器和Literal类型。

解决方案的关键在于使用@overload装饰器定义函数的多个重载版本，每个版本对应一个特定的参数值和返回值类型。 Literal类型确保参数值精确匹配。

以下代码模拟open()函数的行为：

from typing import overload, TypeAlias, Literal

RT: TypeAlias = Literal["r",]
RB: TypeAlias = Literal["rb",]


@overload
def open(
    mode: RT = ...,
) -> str: ...


@overload
def open(
    mode: RB = ...,
) -> bytes: ...


result = open('rb')  # IDE 将识别 result 的类型为 bytes
result = open('r')   # IDE 将识别 result 的类型为 str

代码中，我们定义了两个open函数的重载版本。第一个版本指定mode参数为Literal["r"]，返回值为str；第二个版本指定mode参数为Literal["rb"]，返回值为bytes。 TypeAlias简化了类型定义。 IDE根据mode参数的实际值推断open函数的返回值类型，提供准确的类型提示。

通过@overload和Literal的组合，我们实现了根据参数值动态调整返回值类型提示的功能，使代码更清晰易维护，并充分利用了IDE的类型检查能力。