如何为带可选参数的装饰器添加类型注解？

要让类型检查器正确推断这种「可直接装饰/带参数调用后装饰」的装饰器类型，核心是用**重载（@overload）**定义两种调用场景的签名，结合ParamSpec和TypeVar保留原函数的参数与返回值类型。

具体实现步骤 #

1. 导入必要的类型工具：

import functools
from typing import overload, Callable, ParamSpec, TypeVar

2. 定义参数规格变量P和返回值类型变量R，用来捕获原函数的参数和返回值类型：

P = ParamSpec("P")
R = TypeVar("R")

3. 用@overload定义两种调用场景的类型签名：

   • 场景1：直接用装饰器（不带参数），此时装饰器接收一个函数，返回同类型的函数
   • 场景2：带参数调用装饰器，此时返回一个能接收函数并返回同类型函数的装饰器

4. 实现装饰器本体，并添加对应的类型提示

完整代码示例 #

import functools
from typing import overload, Callable, ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")

@overload
def my_decorator(func: Callable[P, R]) -> Callable[P, R]:
    ...

@overload
def my_decorator(*, param: int = 42) -> Callable[[Callable[P, R]], Callable[P, R]]:
    ...

def my_decorator(func: Callable[P, R] | None = None, *, param: int = 42) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
    if func is None:
        return functools.partial(my_decorator, param=param)
    
    # 这里写装饰器的实际逻辑，比如包装函数
    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"装饰器参数：{param}")
        return func(*args, **kwargs)
    return wrapper

验证两种用法的类型推断 #

# 用法1：直接装饰
@my_decorator
def add(a: int, b: int) -> int:
    return a + b

# 类型检查器能正确推断add的类型是 (int, int) -> int
add(1, 2)  # 正常通过类型检查

# 用法2：带参数装饰
@my_decorator(param=123)
def multiply(x: float, y: float) -> float:
    return x * y

# 类型检查器能正确推断multiply的类型是 (float, float) -> float
multiply(2.5, 4.0)  # 正常通过类型检查

这种写法和Python标准库中functools.lru_cache的类型提示逻辑一致，能让类型检查器准确识别两种使用方式下被装饰函数的类型。

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