CODEWITHSUNDEEP
pythonpython-3.x
Hubro
Hubro

Reputation: 59383

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

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

Upvotes: 5

Views: 1332

Answers (3)

Jonathan Wickens
Jonathan Wickens

Reputation: 769

Since python 3.10 you can now type this correctly using ParamSpec generic with Concatenate to remove arguments from a signature.

The following code type checks correctly in python 3.11 with mypy and 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)

For more information about using ParamSpec to modify call signatures statically see PEP 612

Please note that I removed the dynamic checking/adding of the DB argument resulting in an optional DB argment as this complicates matters further, especially since the ParamSpec does not support keyword arguments.

Upvotes: 2

wim
wim

Reputation: 363083

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

Upvotes: 4

soumya-kole
soumya-kole

Reputation: 1375

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

Upvotes: 2

Related Questions