Python decorator adding argument to function and its signature

user_4172698 0
Viewed 1334

I have multiple classes, many of them having the same initialisation code for one parameter. Therefore I wanted to add the argument with the wrapper.

Since the code is already in production and this parameter is last in all calls, but the signatures have different lengths and the parameter can be position only, it is not trivial to "catch" this argument from the args and kwargs.

The following "works", as long as step is a kwarg, but if not, it is in *args and will be passed to the function, which correctly throws because it got too many arguments:

def stepable(func):
    @functools.wraps(func)
    def wrapper(self, *args, step=1, **kwargs):
        func(self, *args, **kwargs)
        self.step = step  # and other stuff, depending on step
    return wrapper

But even if I would catch it with len(args)>len(inspect.signature(func).parameters) (there are no *args in the functions parameters) the signature shown to the Users is wrong (because I used @wraps).

How can I add the parameter(/default) so that inspect will get it? Or basically "do the inverse of functools.partial"?

python-decorators python
2 Answers

Your problem is that functools.wraps copy the original signature. Here, you will have to manually process and change it. If could be simple enough if you could be sure that none of the wrapped method could have:

  • a step parameter
  • a *args (VAR_POSITIONAL) parameter
  • a **kwargs (VAR_KEYWORD) parameter

And if the step parameter had no default value

But anyway, the inspect module provides everything to deal with signature.

I would define step to be the last POSITIONAL_OR_KEYWORD parameter in the wrapped function

Possible code:

def stepable(func):
    oldsig = inspect.signature(func)
    # search if a VAR_POSITIONAL or VAR_KEYWORD is present
    # if yes insert step parameter before it, else insert it in last position
    params = list(oldsig.parameters.values())
    for i, param in enumerate(params):
        if param.kind == inspect.Parameter.VAR_POSITIONAL:
            break
        if param.kind == inspect.Parameter.VAR_KEYWORD:
            break
    else:
        i = len(params)
    # new parameter name is step or step_[_...] if step if already present
    name = "step"
    while name in oldsig.parameters:
        name += '_'
    newparam = inspect.Parameter(name,
                                 inspect.Parameter.POSITIONAL_OR_KEYWORD,
                                 default = 1)
    params.insert(i, newparam)
    # we can now build the signature for the wrapper function
    sig = oldsig.replace(parameters = params)

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        bound = sig.bind(self, *args, **kwargs) # compute the bound parameter list
        bound.apply_defaults()
        step = bound.arguments[name]      # extract and remove step
        del bound.arguments[name]
        cr = func(*bound.args, **bound.kwargs) # call original function
        self.step = step
        return cr
    wrapper.__signature__ = sig
    return wrapper

Demo:

>>> class A:
    @stepable
    def func(self, a, b=1):
        """This is a test"""
        print(a,b)


>>> a = A()
>>> a.func(5)
5 1
>>> a.step
1
>>> a.func(5,6)
5 6
>>> a.step
1
>>> a.func(5,6,7)
5 6
>>> a.step
7
>>> help(a.func)
Help on method func in module __main__:

func(a, b=1, step=1) method of __main__.A instance
    This is a test
user_35452730

Hope this helps

def stepable(func):
    @functools.wraps(func)
    def wrapper(self, *args, step=1, **kwargs):
        return func(self, *args, **kwargs, step=step)
    return wrapper

class A:
    @stepable
    def __init__(self,a, b, **kwargs):
        self.a = a
        self.b = b
        for k in kwargs:
            setattr(self, k, kwargs[k])

It will pass provided step or 1 to the function Though you have to pass step as named argument only on instantiation and it will make your code more readable.

user_113825670
Related