如何修复函数返回联合类型时mypy的赋值类型不兼容错误？

问题根源 #

当前函数仅声明返回Union[int, str]，mypy无法根据传入参数b的具体值（1或0）推断出对应的返回类型，因此认为赋值给int/str类型变量时存在类型不兼容问题。

最佳方案：使用函数重载（@overload） #

通过typing.overload为函数定义不同输入对应的返回类型，让mypy能精准推断返回值类型，这是类型安全且最符合规范的解决方式。

修复后的代码：

from typing import Union, overload, Literal

@overload
def a(b: Literal[0]) -> str:
    ...

@overload
def a(b: int) -> int:
    ...

def a(b: int) -> Union[int, str]:
    if b:
        return b
    else:
        return '2'

c: int = a(1)  # mypy可推断返回int，无错误
d: str = a(0)  # mypy可推断返回str，无错误

这里用Literal[0]精准限定输入为0时返回字符串，其他整数输入返回整数，让类型检查逻辑更严谨。

其他可选方案（不推荐） #

• 类型断言（cast）：手动指定返回类型，但需要开发者自行保证逻辑正确性，类型检查器无法验证，存在风险：

  from typing import Union, cast
  
  def a(b: int) -> Union[int, str]:
      if b:
          return b
      else:
          return '2'
  
  c: int = cast(int, a(1))
  d: str = cast(str, a(0))
  

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