如何为返回装饰器、上下文管理器或自定义对象的函数添加类型注解

我开发了一款可为测试添加功能的库，支持三种使用方式：

def test() -> None:
    with custom_functionality():
        print("test")

@custom_functionality
def test_deco() -> None:
    print("test_deco")


def test_manual() -> None:
    custom_func = custom_functionality()
    custom_func.start()
    print("test_manual")
    custom_func.stop()

该库的内部实现逻辑如下：

from typing import ContextManager, Any, Callable, Optional, TypeVar, Union
from typing_extensions import ParamSpec

P = ParamSpec('P')
T = TypeVar('T')


class ContextDecorator:
    def __init__(self, func: Optional[Callable[P, T]] = None) -> None:
        self.func = func

    def __call__(self) -> Union[Callable[P, T], T]:
        self.start()
        x = self.func()
        self.stop()
        return x

    def start(self) -> None:
        self.__enter__()

    def stop(self) -> None:
        self.__exit__()

    def __enter__(self) -> None:
        pass

    def __exit__(self, *args: Any) -> None:
        pass


def custom_functionality(func: Optional[Callable[P, T]] = None):
    if func:
        def wrapper() -> Callable[P, T]:
            context_deco = ContextDecorator(func)
            return context_deco()

        return wrapper
    else:
        return ContextDecorator()

问题 #

需要为custom_functionality函数添加返回类型注解，但直接使用三种类型的Union无法让MyPy区分不同场景下的类型：

• 标准装饰器场景（@custom_functionality）
• 上下文管理器/手动调用start/stop场景（custom_functionality()）

解决方案：使用@overload定义多场景类型签名 #

Python的typing模块提供了@overload装饰器，可以为同一个函数定义多个类型签名，类型检查器会根据实际调用方式匹配对应的签名，从而精确推断返回类型。

修改后的代码（添加类型注解） #

from typing import ContextManager, Any, Callable, Optional, TypeVar, Union, overload
from typing_extensions import ParamSpec

P = ParamSpec('P')
T = TypeVar('T')


class ContextDecorator:
    def __init__(self, func: Optional[Callable[P, T]] = None) -> None:
        self.func = func

    # 修正__call__的类型注解：适配带括号的装饰器场景，接收func并返回包装后的函数
    def __call__(self, func: Callable[P, T]) -> Callable[P, T]:
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            self.start()
            result = func(*args, **kwargs)
            self.stop()
            return result
        return wrapper

    def start(self) -> None:
        self.__enter__()

    def stop(self) -> None:
        self.__exit__()

    def __enter__(self) -> None:
        pass

    def __exit__(self, *args: Any) -> None:
        pass


# 重载1：传入func参数时，作为无括号装饰器返回包装后的函数
@overload
def custom_functionality(func: Callable[P, T]) -> Callable[P, T]:
    ...

# 重载2：不传入参数时，返回ContextDecorator实例，支持上下文管理器、手动start/stop、带括号装饰器
@overload
def custom_functionality() -> ContextDecorator:
    ...

# 实际实现的类型注解
def custom_functionality(func: Optional[Callable[P, T]] = None) -> Union[Callable[P, T], ContextDecorator]:
    if func:
        # 修正wrapper的返回类型：与原函数返回类型保持一致
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            context_deco = ContextDecorator()
            context_deco.start()
            result = func(*args, **kwargs)
            context_deco.stop()
            return result
        return wrapper
    else:
        return ContextDecorator()

说明 #

1. @overload的作用：为custom_functionality定义了两种明确的调用场景类型：
   • 传入func参数时（如@custom_functionality装饰函数），类型检查器匹配第一个重载，明确返回与原函数同类型的Callable[P, T]。
   • 不传入参数时（如custom_functionality()），匹配第二个重载，返回ContextDecorator实例，此时类型检查器能正确识别该实例的start()、stop()方法以及上下文管理器协议方法。
2. 修正原有代码的类型问题：
   • 调整ContextDecorator.__call__的类型注解，使其符合标准装饰器类的行为，支持@custom_functionality()这类带括号的装饰器用法。
   • 修正wrapper函数的返回类型，确保与原函数的返回类型一致，避免类型不匹配。

这样修改后，MyPy等类型检查工具就能在不同使用场景下正确推断类型，解决了Union无法区分场景的问题。

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