How to decorate a function in python in a way that adds a new argument and creates the correct signature?

user_4135063 0
Viewed 72

I have a decorator that "wraps" a function in a new one that has one extra positional argument in the beginning:

import functools

def my_decorator(f):
    @functools.wraps(f)
    def wrapped(message: str, *args, **kwargs):
        print('Decorated:', message)
        return f(*args, **kwargs)
    return wrapped

Here is an example on how my decorator is used:

@my_decorator
def foo(a: int, b: str) -> None:
    print('Hello world', [a, b])

foo('test', 2, '3')
# Decorated: test
# Hello world [2, '3']

So far so good. However, I'd like my decorator to create a new function with the correct signature. But since I used @functools.wraps, I am getting the same signature as the original function:

import inspect

print('Signature:', inspect.signature(foo))
# Signature: (a: int, b: str) -> None

But I want to get (message: str, a: int, b: str) -> None instead. How to do this?

I need the correct signature to get correct type checks and also because I am using a framework that relies on function signatures for its behavior.

python-typing python-decorators python
2 Answers

This isn't the most elegant solution since it's set manually not automatically, but something like this could work. Not sure about automatically adjusting it based on the signature of the wrapped command itself.

import functools
from inspect import Parameter, signature

def my_decorator(f):

    @functools.wraps(f)
    def wrapped(message: str, *args, **kwargs):
        print('Decorated:', message)
        return f(*args, **kwargs)

    message_param = Parameter('message', Parameter.POSITIONAL_OR_KEYWORD, annotation=str)
    sig = signature(f)
    parameters = [message_param] + list(sig.parameters.values())
    wrapped.__signature__ = sig.replace(parameters=parameters)

    return wrapped
user_46902370

Regarding the type annotations, this will be the complete and correct one:

import functools
from typing import Callable, ParamSpec, Protocol, TypeVar

P = ParamSpec("P")
RT = TypeVar("RT")
T_co = TypeVar("T_co", covariant=True)


class WithMessage(Protocol[P, T_co]):
    def __call__(self, message: str, *args: P.args, **kwds: P.kwargs) -> T_co:
        ...


def my_decorator(f: Callable[P, RT]) -> WithMessage[P, RT]:
    @functools.wraps(f)
    def wrapped(message: str, *args: P.args, **kwargs: P.kwargs) -> RT:
        print("Decorated:", message)
        return f(*args, **kwargs)

    return wrapped


@my_decorator
def foo(a: int, b: str) -> None:
    print("Hello world", [a, b])

Unfortunately, inspect.signature does not seem to respect even the complete and correct type annotations:

# prints: "Signature: (a: int, b: str) -> None"
# instead of "Signature: (message: str, a: int, b: "str) -> None
print("Signature:", inspect.signature(foo))

Here's a playground with this solution.

user_80910930
Related