Preserve exact function signature in chain of decorators using functools.wraps

Question

In Python 3.4+, functools.wraps preserves the signature of the function it wraps. Unfortunately, if you create decorators that are meant to be stacked on top of each other, the second (or later) decorator in the sequence will be seeing the generic *args and **kwargs signature of the wrapper and not preserving the signature of the original function all the way at the bottom of the sequence of decorators. Here's an example.

from functools import wraps    

def validate_x(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        assert kwargs['x'] <= 2
        return func(*args, **kwargs)
    return wrapper

def validate_y(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        assert kwargs['y'] >= 2
        return func(*args, **kwargs)
    return wrapper

@validate_x
@validate_y
def foo(x=1, y=3):
    print(x + y)


# call the double wrapped function.
foo()

This gives

-------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-5-69c17467332d> in <module>
     22
     23
---> 24 foo()

<ipython-input-5-69c17467332d> in wrapper(*args, **kwargs)
      4     @wraps(func)
      5     def wrapper(*args, **kwargs):
----> 6         assert kwargs['x'] <= 2
      7         return func(*args, **kwargs)
      8     return wrapper

KeyError: 'x'

and if you switch the order of the decorators, you get the same key error for 'y'.

I tried replacing wraps(func) with wraps(func.__wrapped__) in the second decorator, but this still doesn't work (not to mention it requires the programmer to explicitly know where in the stack of decorators they are working for given wrapper functionality).

I also took a look at inspect.signature(foo) and this seems to give the right thing, but I found that this is because inspect.signature has a follow_wrapped parameter that defaults to True so it somehow knows to follow the sequence of wrapped functions, but apparently the regular method call framework for invoking foo() will not follow this same protocol for resolve args and kwargs of the outer decorated wrapper.

How can I just have wraps faithfully passthrough the signature so that wraps(wraps(wraps(wraps(f)))) (so to speak) always faithfully replicated the signature of f?

score 2 · Answer 1 · answered Jan 22 '20 at 15:44

2

You are not actually passing any arguments to you function foo so *args and **kwargs are empty for both decorators. If you pass arguments the decorators will work just fine

foo(x=2, y = 3) # prints 5

You can try to get default function arguments using inspect

answered Jan 22 '20 at 15:44

hurlenko

1,363
2
12
17

3

I don't know who downvoted this but it's totally correct and totally relevant - even if it isn't enough to solve the issue. – bruno desthuilliers Jan 22 '20 at 16:15

JBirdVegas · Answer 2 · 2020-01-24T14:49:10.910

You can't really get the default values without using inspect and you also need to account for positional args (*args) vs keyword args (**kwargs). So normalize the data if it's there if it's missing then inspect the function

import inspect
from functools import wraps


def get_default_args(func):
    signature = inspect.signature(func)
    return {
        k: v.default
        for k, v in signature.parameters.items()
        if v.default is not inspect.Parameter.empty
    }


def validate_x(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if args and not kwargs and len(args) == 2:
            kwargs['x'] = args[0]
            kwargs['y'] = args[1]
            args = []
        if not args and not kwargs:
            kwargs = get_default_args(func)
        assert kwargs['x'] <= 2
        return func(*args, **kwargs)

    return wrapper


def validate_y(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if args and not kwargs and len(args) == 2:
            kwargs['x'] = args[0]
            kwargs['y'] = args[1]
            args = []
        if not args and not kwargs:
            kwargs = get_default_args(func)
        assert kwargs['y'] >= 2
        return func(*args, **kwargs)

    return wrapper


@validate_x
@validate_y
def foo(x=1, y=3):
    print(x + y)


# call the double wrapped function.
foo()
# call with positional args
foo(1, 4)
# call with keyword args
foo(x=2, y=10)

This prints

4
5
12

Note that the interior decorator, such as `@validate_y` works just fine in my example. It is using `kwargs` from the wrapped function without any `inspect` manipulation. The error only happens for the outer decorator that `wraps(...)` another decorated function. — ely, Jan 22 '20 at 15:52
When I ran your code in the debugger both the args and kwargs were empty. This is because the default values are filled in at the time the function is called not before. That's why I added the normalization of `args` -> `kwargs` and needed `inspect` to get the default values. — JBirdVegas, Jan 22 '20 at 15:54
That is true for the outer decorator, but doesn't seem to be true for the inner decorator. — ely, Jan 22 '20 at 17:05
@ely The interior decorator `@validate_y` also doesn't work in your example; it's just that you didn't give it a chance to fail, because `@validate_x` failed first. The behaviour is the same whether you use one or both decorators; either way it raises a `KeyError`. — kaya3, Jan 22 '20 at 18:53

kaya3 · Answer 3 · 2020-01-22T18:49:25.793

0

Your diagnosis is incorrect; in fact, functools.wraps preserves the signature of the double-decorated function:

>>> import inspect
>>> inspect.signature(foo)
<Signature (x=1, y=3)>

We can also observe that it is not a problem with calling the function with the wrong signature, since that would raise a TypeError, not a KeyError.

You seem to be under the impression that when just one decorator is used, kwargs will be populated with the argument default values. This doesn't happen at all:

def test_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print('args:', args)
        print('kwargs:', kwargs)
        return func(*args, **kwargs)
    return wrapper

@test_decorator
def foo(x=1):
    print('x:', x)

The output is:

>>> foo()
args: ()
kwargs: {}
x: 1

So as you can see, neither args nor kwargs receives the argument's default value, even when just one decorator is used. They are both empty, because foo() calls the wrapper function with no positional arguments and no keyword arguments.

The actual problem is that your code has a logical error. The decorators validate_x and validate_y expect the arguments to be passed as keyword arguments, but in fact they might be passed as positional arguments or not at all (so the default values would apply), in which case 'x' and/or 'y' won't be present in kwargs.

There is no easy way to make your decorators work with an argument which could be passed as either keyword or positional; if you make the arguments keyword-only, then you can test whether 'x' or 'y' are in kwargs before validating them.

def validate_x(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if 'x' in kwargs and kwargs['x'] > 2:
            raise ValueError('Invalid x, should be <= 2, was ' + str(x))
        return func(*args, **kwargs)
    return wrapper

@validate_x
def bar(*, x=1): # keyword-only arg, prevent passing as positional arg
    ...

It's usually better to explicitly raise an error, instead of using assert, because your program can be run with assert disabled.

Beware also that it's possible to declare a function like @validate_x def baz(*, x=5): ... where the default x is invalid. This won't raise any error because the default argument value isn't checked by the decorator.

edited Jan 22 '20 at 18:49

answered Jan 22 '20 at 15:55

kaya3

47,440
4
68
97

I addressed this in my question already. Note that `inspect.signature` has a `follow_wrapped` parameter that is True by default, but the actual signature getting used in the function call machinery is different, not following the wraps. – ely Jan 22 '20 at 15:56
The actual signature of the wrapper is `(*args, **kwargs)` because that's how you declared it. As I said, the error is not in the signature, but in the logic of your decorators. You are unconditionally accessing `kwargs['x']` when `kwargs` is not guaranteed to contain `'x'` as a key. If the signature were the problem then you'd get a `TypeError` for calling `foo` with the wrong arguments, not a `KeyError` from the body of the function; the fact that the function's body executes at all means the arguments are acceptable according to the signature. – kaya3 Jan 22 '20 at 15:59
For an analogy, if you defined `def wrapper(*args, **kwargs): raise ValueError()` then `foo()` would raise a ValueError because that's the declared behaviour of the wrapper function; the fact that calling the decorated function raises an error doesn't mean its signature is wrong. – kaya3 Jan 22 '20 at 16:02
But in the interior decorator, references to plain `kwargs` are able to access the default passed to the decorated function, with no use of `inspect`. It's only the outer decorator that fails. – ely Jan 22 '20 at 17:06
`kwargs` would be populated with the wrapped function's defaults, if the wrapper function is called with no args. E.g. if `wrapper` wraps `func` using `functools.wraps` and `func` defines keyword args with defaults, then it should be impossible for `wrapper` to ever be called with empty `kwargs`. It would _always_ have the ones coming from the function it wraps, at minimum, unavoidably. – ely Jan 22 '20 at 17:14
*"But in the interior decorator, references to plain kwargs are able to access the default passed to the decorated function, with no use of inspect. It's only the outer decorator that fails."* That's just not true; if you use just the `@validate_y` decorator without `@validate_x` then it fails with `KeyError: 'y'`. The only reason the inner decorator doesn't raise an error when you use both decorators, is because the outer decorator raises an error before it gets a chance to call the inner wrapped function. – kaya3 Jan 22 '20 at 18:34
I have edited to add a demonstration that `kwargs` is not, in fact, populated with the argument defaults when you have just one decorator. – kaya3 Jan 22 '20 at 18:40

Preserve exact function signature in chain of decorators using functools.wraps

3 Answers3