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

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.

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

为此我创建了一个装饰器:

def provide_db(fn):
    ...

这个装饰器检查关键字参数是否包含键“db”，如果没有，它会创建一个连接并将其传递给函数。用法:

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

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

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

目前使用 Python 的类型提示是否可行？如果是这样，它是如何完成的？

---

这是 provide_db 函数:

def provide_db(fn):
    """Decorator that defaults the db argument to a new connection

    This pattern allows callers to not have to worry about passing a db
    parameter, but also allows functions to pass db connections to each other to
    avoid the overhead of creating unnecessary connections.
    """

    if not "db" in fn.__code__.co_varnames:
        raise ValueError("Wrapped function doesn't accept a db argument")

    db_arg_index = fn.__code__.co_varnames.index("db")

    @wraps(fn)
    def wrapper(*args, **kwargs) -> Result:
        if len(args) > db_arg_index and args[db_arg_index] is not None:
            pass  # db was passed as a positional argument
        elif "db" in kwargs and kwargs["db"] is not None:
            pass  # db was passed as a keyword argument
        else:
            kwargs["db"] = connect()

        return fn(*args, **kwargs)

    wrapper.__annotations__ = Optional[fn.__annotations__["db"]]

    return wrapper

Answer 1

函数注释在 the datamodel 中记录为可写在PEP 526 .

按照这个简化的例子:

from __future__ import annotations
from typing import Optional

def provide_db(func):
    func.__annotations__["db"] = Optional[func.__annotations__["db"]]
    return func

@provide_db
def f(db: Connection):
    pass

print(f.__annotations__)