Can decorators modify type hints in the function signature of wrapped functions?

user_388916 0
Viewed 398

I have a bunch of functions that look like this:

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.

What these functions all have in common is that they take a Connection parameter called db that they use to modify a database. For performance reasons, I want the functions to be able to pass the db parameter between each other and not create a new connection every time. However, for convenience reasons, I also don't want to have to create and pass the db parameter every time I call the functions myself.

For that reason I have created a decorator:

def provide_db(fn):
    ...

This decorator checks if the keyword arguments contain the key "db", and if not, it creates a connection and passes it to the function. Usage:

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

This works perfectly! I can now call the database functions without worrying about connecting to the database, and the functions can pass the db parameters to each other.

However, for this to be typed properly, the decorator needs to modify the function signature of the wrapped function, changing the db parameter from Connection to Optional[Connection].

Is this currently possible with Python's type hints? If so, how is it done?

This is the provide_db function:

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
python-3.x python
2 Answers

Function annotations are documented as writable in the datamodel and in PEP 526.

Follow this simplified example:

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__)
user_6740390

For this requirement, it seems be good to have all these functions in a class.

class DbOps(object):
    def __init__(self, db):
        self.__db = db

    def insert_user(user: User) -> None:
        #use self__.db to connect db and insert user

Then :

db_ops = DbOps(db)
db_ops.insert_user(...)
etc
user_21225020
Related