装饰器可以修改包装函数的函数签名中的类型提示吗？

Question

我有很多类似这样的函数：

def insert_user(user: User, db: Connection) -> None:
    ...

def add_role_to_user(user: User, role: str, db: Connection) -> None:
    ...

def create_post(post: Post, owner: User, db: Connection) -> None:
    ...

# etc.

这些函数的共同点是它们都采用一个Connection名为 的参数db，用于修改数据库。出于性能原因，我希望函数能够db在彼此之间传递参数，而不是每次都创建新连接。db但是，为了方便起见，我也不希望每次自己调用函数时都必须创建并传递参数。

因此我创建了一个装饰器：

def provide_db(fn):
    ...

该装饰器检查关键字参数是否包含键“db”，如果不包含，则创建一个连接并将其传递给函数。用法：

@provide_db
def insert_user(user: User, db: Connection) -> None:
    ...

这非常有效！我现在可以调用数据库函数，而不必担心连接到数据库，并且函数可以相互传递数据库参数。

但是，为了正确键入此内容，装饰器需要修改包装函数的函数签名，将参数db从更改Connection为Optional[Connection]。

目前，Python 的类型提示可以实现这一点吗？如果是这样，是如何做到的？

---

这是provide_db函数：

def insert_user(user: User, db: Connection) -> None:
    ...

def add_role_to_user(user: User, role: str, db: Connection) -> None:
    ...

def create_post(post: Post, owner: User, db: Connection) -> None:
    ...

# etc.

Answer 1

ParamSpec从 python 3.10 开始，您现在可以使用泛型 with正确键入此内容，Concatenate以从签名中删除参数。

以下代码类型在 python 3.11 中使用 mypy 和 pylance 进行正确检查。

from typing import Callable, ParamSpec, TypeVar, Concatenate
from functools import wraps

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

class DB:
    pass

db = DB()

def provide_db(fn: Callable[Concatenate[DB, P], R]) -> Callable[P, R]:
    @wraps(fn)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        return fn(db, *args, **kwargs)

    return wrapper

@provide_db
def foo(db: DB, x: int) -> int:
    return x

foo(1)

有关使用ParamSpec静态修改调用签名的更多信息，请参阅PEP 612

请注意，我删除了数据库参数的动态检查/添加，从而产生了可选的数据库参数，因为这使问题进一步复杂化，特别是因为 ParamSpec 不支持关键字参数。