Python 注释不同输入情况的返回类型

Question

我有一个函数，它可以在给定 bool 输入参数的情况下返回不同的类型。目前，我只是使用 union 对其进行注释:

def fn(x: int, ret_str: bool) -> Union[int, str]:
  if ret_str:
    return str(x)
  else:
    return x

但是，如果我这样注释，该函数的输出也将被键入为联合类型。

out = fn(8, ret_str=False)  # out will have type Union[int, str] but we know it's int

有没有办法对其进行注释，以便类型检查器知道我是否使用文字 False 调用函数，返回类型必须是 int 而不是str?

Answer 1

是的，这可以使用typing 模块和overload/Literal。它不是最优雅的解决方案，因为它非常冗长(与大多数类型化语言相比)，但会导致类型检查器根据传递的参数关联返回类型。在库/模块之外的大多数情况下，我可能会避免这样做，或者我知道会从外部源重复引用的其他代码。

from typing import Literal, overload


@overload
def fn(x: int, ret_str: Literal[False]) -> int:
    ...


@overload
def fn(x: int, ret_str: Literal[True]) -> str:
    ...


def fn(x: int, ret_str: bool) -> int | str:
    if ret_str:
        return str(x)
    return x

如果想要像这样实现更高级的类型检查，将此信息存储在单独的 pyi 文件中而不是直接在您的程序中通常是有益的。

最后，请注意 @overload 装饰器纯粹用于类型检查器。你不能用它把不同的功能放到一个函数中，因为函数的最终声明总是会覆盖任何以前的定义(不管是否使用装饰器)。因此，在重载函数定义中放置省略号 (...) 是范例。