如何为参数为0/其他整数时返回不同类型的函数添加重载注解？

问题背景 #

需要实现一个函数：参数为0时返回MyObjectA实例，参数为其他整数时返回MyObjectB实例。现有实现无法让Mypy正确推断分支类型，且尝试重载时遇到签名重叠的问题。

现有实现与错误 #

基础函数实现：

def foo(val: int) -> MyObjectA | MyObjectB:
    if val == 0:
        return MyObjectA()
    return MyObjectB()

执行以下赋值时，Mypy触发类型不兼容错误：

baz: MyObjectA = foo(0)
bar: MyObjectB = foo(1)

错误提示：

Incompatible types in assignment (expression has type "MyObjectA | MyObjectB", variable has type "MyObjectA") [assignment]

尝试重载时的问题 #

使用Literal[0]匹配参数为0的情况，但无法直接表示“除0外的所有整数”：

@overload
def foo(val: Literal[0]) -> MyObjectA:
    ...

# 无法枚举所有非0整数
@overload
def foo(val: Literal[1,2,3,4,5,6, ...]) -> MyObjectB: 
    ...

调整重载顺序后，Mypy提示签名重叠错误（已翻译为中文）：

@overload
def foo(val: Literal[0]) -> MyObjectA:
    ...

@overload
def foo(val: int) -> MyObjectB:
    ...

错误提示：

重载函数签名1和2存在重叠，且返回类型不兼容 [misc]

解决方案 #

方案1：使用Not类型排除特定字面量 #

Python 3.11+ 可直接使用标准库typing.Not，Python<3.11则需安装typing_extensions并导入Not。通过int & Not[Literal[0]]明确表示“除0外的所有整数”，避免签名重叠：

from typing import overload, Literal
# Python 3.11+ 导入方式
from typing import Not
# Python <3.11 导入方式
# from typing_extensions import Not

class MyObjectA:
    pass

class MyObjectB:
    pass

@overload
def foo(val: Literal[0]) -> MyObjectA:
    ...

@overload
def foo(val: int & Not[Literal[0]]) -> MyObjectB:
    ...

def foo(val: int) -> MyObjectA | MyObjectB:
    if val == 0:
        return MyObjectA()
    return MyObjectB()

此时执行目标赋值代码，Mypy会正确推断类型：

baz: MyObjectA = foo(0)  # 推断为MyObjectA，无错误
bar: MyObjectB = foo(1)  # 推断为MyObjectB，无错误

方案2：使用TypeVar结合约束 #

若无法使用Not类型，可通过绑定整数的TypeVar实现重载：

from typing import overload, Literal, TypeVar

class MyObjectA:
    pass

class MyObjectB:
    pass

IntT = TypeVar("IntT", bound=int)

@overload
def foo(val: Literal[0]) -> MyObjectA:
    ...

@overload
def foo(val: IntT) -> MyObjectB:
    ...

def foo(val: int) -> MyObjectA | MyObjectB:
    if val == 0:
        return MyObjectA()
    return MyObjectB()

该方案类型推断精度略低于方案1，但可兼容更早的Python版本。

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