Python decorator that returns a function with one or more arguments replaced

时间:2016-10-15 17:13:15

标签: python decorator partial python-decorators

I would like to create a decorator to a set of functions that replaces one or more of the arguments of the functions. The first thought that came to my mind was to create a decorator that returns a partial of the function with the replaced arguments. I'm unhappy with the way the decorated function is called, but even when it's being called "properly", i get a TypeError.

Here is some example code:

def decor(func, *args, **kwargs):
    def _new_func(*args, **kwargs):
        return partial(func, *args, **kwargs, v=100)
    return _new_func

@decor
def some_function(a, v, c):
    return a, v, c

some_function(1,2,3)  # functools.partial(<function some_function at 0x7f89a8bed8c8>, 1, 2, 3, v=100)
some_function(1,2,3)(1,2,3)  # TypeError: some_function() got multiple values for argument 'v'

I'm sure there is an easy way to create a decorator that replaces some of the arguments, but haven't figured it out yet.

1 个答案:

答案 0 :(得分:3)

You'd have to return the partial as the decoration result:

def decor(func):
    return partial(func, v=100)

However, this always sets v=100, even if you passed in another value for v by position. You'd still have the same issue.

You'd need to create a decorator that knows what positional argument v is, and look for it there and as a keyword argument:

from inspect import getargspec

def decor(func):
    vpos = getargspec(func).args.index('v')
    def wrapper(*args, **kwargs):
        if len(args) > vpos:
            args = list(args)
            args[vpos] = 100
        else:
            kwargs['v'] = 100
        return func(*args, **kwargs)        
    return wrapper

The above decorator will set the v argument to 100, always. Wether you tried to set v as a positional argument or as a keyword argument, in the end it'll be set to 100 anyway:

>>> @decor
... def some_function(a, v, c):
...     return a, v, c
...
>>> some_function(1, 2, 3)
(1, 100, 3)
>>> some_function(1, v=2, c=3)
(1, 100, 3)

If you only wanted to provide a default argument for v if it wasn't explicitly set, you'd have to invert the tests:

def decor(func):
    vpos = getargspec(func).args.index('v')
    def wrapper(*args, **kwargs):
        if len(args) <= vpos and 'v' not in kwargs:
            kwargs['v'] = 100
        return func(*args, **kwargs)        
    return wrapper

at which point v is only provided if not already set:

>>> @decor
... def some_function(a, v, c):
...     return a, v, c
...
>>> some_function(1, c=3)  # v not set
(1, 100, 3)
>>> some_function(1, 2, c=3)  # v is set
(1, 2, 3)
>>> some_function(1, v=2, c=3)  # v is set as keyword argument
(1, 2, 3)