如何为装饰器添加类型提示，将指定关键字参数设为可选

针对你遇到的将ConnPool参数转为可选、同时保留完整类型提示的需求，可以通过**函数重载（@overload）**结合functools.wraps实现，配合ParamSpec和TypeVar保留原函数的参数与返回值类型，完美解决类型提示问题。

完整实现代码 #

from typing import Callable, List, Optional, ParamSpec, TypeVar, Awaitable, overload
from functools import wraps

# 示例ConnPool类（根据实际业务逻辑替换）
class ConnPool:
    async def close(self) -> None:
        """模拟连接池资源释放逻辑"""
        pass

# 定义参数和返回值类型变量
P = ParamSpec("P")
R = TypeVar("R")

# 重载装饰器签名，为IDE提供两种合法调用的类型提示
@overload
def auto_pool(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
    ...

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

def auto_pool(func: Callable[P, Awaitable[R]]) -> Callable[..., Awaitable[R]]:
    @wraps(func)
    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        # 提取传入的连接池，或标记需要创建新池
        pool: Optional[ConnPool] = kwargs.pop("pool", None)
        created_pool = False
        
        try:
            if pool is None:
                # 动态创建新连接池的逻辑
                pool = ConnPool()
                created_pool = True
            # 调用原函数并传入连接池
            return await func(*args, pool=pool, **kwargs)
        finally:
            # 仅关闭当前装饰器创建的连接池，避免影响外部传入的复用池
            if created_pool and pool is not None:
                await pool.close()
    
    return wrapper

使用示例 #

@auto_pool
async def list_jobs(namespaces: List[str], *, pool: ConnPool) -> List[str]:
    # 业务逻辑示例
    return [f"{namespace}:sample_job" for namespace in namespaces]

方案优势 #

1. 精准类型提示：IDE会自动识别两种合法调用方式：
   • 传入已有连接池：await list_jobs(["prod", "dev"], pool=my_pool)
   • 不传入连接池：await list_jobs(["prod", "dev"])
2. 完整保留原函数签名：位置参数、关键字-only参数的类型约束全部生效，不会破坏原函数的参数校验逻辑
3. 资源安全：自动管理动态创建的连接池生命周期，避免资源泄漏

内容的提问来源于stack exchange，提问作者Christian P.