如何使用Mypy为参数重叠的重载方法添加类型提示？

简化示例的类型错误 #

我遇到一个Mypy类型提示问题，代码如下：

from typing import overload, Union

@overload
def unsafe_func(x: int) -> int: ...

@overload
def unsafe_func(x: object) -> str: ...

def unsafe_func(x: object) -> Union[int, str]:
    if isinstance(x, int):
        return 42
    else:
        return "some string"

Mypy报错：

main.py:4: error: Overloaded function signatures 1 and 2 overlap with incompatible return types
Found 1 error in 1 file (checked 1 source file)

我知道原因是int是object的子类，导致unsafe_func(42)的调用无法匹配明确的重载签名，但不知道怎么修改类型提示，来准确表达“传入int返回int，传入其他object类型返回str”的逻辑。

实际应用场景：带可选参数的装饰器 #

这个问题来自我要实现的一个带可选参数的装饰器，代码如下：

from inspect import isclass

def catcher(function_or_exception):
    # 装饰器带参数的情况：传入的是异常类
    if isclass(function_or_exception) and issubclass(function_or_exception, BaseException):
        def decorator(decorated):
            def decorate():
                try:
                    decorated()
                except function_or_exception:
                    print("An error occurred")
            return decorate
        return decorator

    # 装饰器不带参数的情况：传入的是函数
    else:
        def decorate():
            try:
                function_or_exception()
            except Exception:
                print("An error occurred")
        return decorate

@catcher
def my_func_1():
    "1" + 1

@catcher(ZeroDivisionError)
def my_func_2():
    1 / 0


my_func_1()
my_func_2()

修复简化示例的重载冲突 #

要解决重载签名重叠的问题，核心是让第二个重载的参数类型明确排除int。可以用以下两种方法：

方法1：使用Mypy的补集类型（Python 3.10+） #

利用Mypy支持的~操作符表示类型补集，结合TypeVar定义排除int的object子类：

from typing import overload, TypeVar, Union

NonInt = TypeVar("NonInt", bound=~int)

@overload
def unsafe_func(x: int) -> int: ...

@overload
def unsafe_func(x: NonInt) -> str: ...

def unsafe_func(x: object) -> Union[int, str]:
    if isinstance(x, int):
        return 42
    else:
        return "some string"

方法2：用Protocol定义非int类型 #

通过Protocol自定义一个只匹配非int对象的类型：

from typing import overload, Protocol, Union

class NonInt(Protocol):
    @classmethod
    def __instancecheck__(cls, instance: object) -> bool:
        return not isinstance(instance, int)

@overload
def unsafe_func(x: int) -> int: ...

@overload
def unsafe_func(x: NonInt) -> str: ...

def unsafe_func(x: object) -> Union[int, str]:
    if isinstance(x, int):
        return 42
    else:
        return "some string"

这两种方式都能让Mypy正确区分调用场景：传入int时匹配第一个重载，返回int；传入其他object子类时匹配第二个重载，返回str。

修复装饰器的类型提示 #

针对带可选参数的装饰器，需要用重载区分两种调用模式，同时用TypeVar保留被装饰函数的类型信息：

from inspect import isclass
from typing import overload, Callable, TypeVar, Type, BaseException

# 定义TypeVar，用于保留被装饰函数的类型
F = TypeVar("F", bound=Callable[..., None])
# 定义异常类的TypeVar
E = TypeVar("E", bound=Type[BaseException])

@overload
def catcher(func: F) -> F: ...

@overload
def catcher(exc_type: E) -> Callable[[F], F]: ...

def catcher(function_or_exception: F | E) -> F | Callable[[F], F]:
    if isclass(function_or_exception) and issubclass(function_or_exception, BaseException):
        def decorator(decorated: F) -> F:
            def decorate() -> None:
                try:
                    decorated()
                except function_or_exception:
                    print("An error occurred")
            return decorate  # type: ignore[return-value]
        return decorator

    else:
        func: F = function_or_exception
        def decorate() -> None:
            try:
                func()
            except Exception:
                print("An error occurred")
        return decorate  # type: ignore[return-value]

@catcher
def my_func_1() -> None:
    "1" + 1

@catcher(ZeroDivisionError)
def my_func_2() -> None:
    1 / 0

my_func_1()
my_func_2()

说明：

• 两个重载分别对应装饰器的两种使用方式：无参数时直接接收函数并返回同类型函数；带参数时接收异常类，返回一个接收函数并返回函数的装饰器。
• type: ignore[return-value]是因为内部嵌套的decorate函数当前只支持无参无返回的函数，如果需要支持带参数或有返回值的函数，需要进一步调整内部函数的类型定义（比如用*args, **kwargs传递参数）。

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