python: argument processing

when writing a function wrapper, we often need to play with its arguments; there are some handy utilities in python that can help;

we show how to use these tools with a simple example; we wrap function:

def foo(a0, a1=1, a2=2, a3=3):
    print(a0, a1, a2, a3)

so that it keeps the same signature but always runs with a2 = 102 no matter what arguments were passed in;

note that the solution provided in this article does not work if foo is a built-in function defined in c;

solution

the solution is really neat:

from functools import wraps
from inspect import signature

@wraps(foo)
def bar(*args, **kwargs):
    sig = signature(foo)
    ba = sig.bind(*args, **kwargs)
    ba.arguments['a2'] = 102        #   set `a2` to `102`;
    foo(*ba.args, **ba.kwargs)

here bar is a wrapper of foo with exactly the same signature:

>>> foo(200, a2=202, a3=203)
200 1 202 203
>>> bar(200, a2=202, a3=203)
200 1 102 203

analysis

the solution has these steps:

get the function signature, which includes parameters;
```
>>> sig = signature(foo)
```
signature returns the call signature of a function; for each parameter of the function, the signature stores a Parameter object in its parameters collection (an ordered dict keyed by parameter names);
```
>>> print(sig.parameters)
OrderedDict([('a0', <Parameter "a0">), ..., ('a3', <Parameter "a3=3">)])
```
bind arguments to parameters;
```
>>> ba = sig.bind(*args, **kwargs)
```
sig.bind creates a mapping from positional and keyword arguments to parameters; the return type is BoundArguments, which has an arguments collection (an ordered dict keyed by argument names) containing explicitly bound arguments;
```
>>> print(ba.arguments)
OrderedDict([('a0', 200), ('a2', 202), ('a3', 203)])
```

update bound arguments;

>>> ba.arguments['a2'] = 102

this is just a dict update;

>>> print(ba.arguments)
OrderedDict([('a0', 200), ('a2', 102), ('a3', 203)])

call the wrapped function with updated bound arguments;
```
>>> foo(*ba.args, **ba.kwargs)
```
ba.args and ba.kwargs are dynamically computed from the ba.arguments collection; therefore, changes in ba.arguments are reflected in ba.args and ba.kwargs, with which we can call the wrapped function;

caveat

the above solution works even if we omit @wraps in our wrapper; but it does not work when the wrapped function has already been wrapped without @wraps; for example:

def baz(*args, **kwargs):
    foo(*args, **kwargs)

we will not be able to wrap baz using the same procedure that wraps foo, because a2 is not a parameter of baz;

from functools import wraps
from inspect import signature

@wraps(baz)
def bar(*args, **kwargs):
    sig = signature(baz)
    ba = sig.bind(*args, **kwargs)
    ba.arguments['a2'] = 102        #   set `a2` to `102`;
    baz(*ba.args, **ba.kwargs)

output:

>>> foo(200, a2=202, a3=203)
200 1 202 203
>>> baz(200, a2=202, a3=203)
200 1 202 203
>>> bar(200, a2=202, a3=203)
200 1 202 203                       #   wrong: `202` not `102`;

to fix this problem, add @wraps:

from functools import wraps

@wraps(foo)
def baz(*args, **kwargs):
    foo(*args, **kwargs)

now we can wrap baz using the same procedure as above;

output:

>>> foo(200, a2=202, a3=203)
200 1 202 203
>>> baz(200, a2=202, a3=203)
200 1 202 203
>>> bar(200, a2=202, a3=203)
200 1 102 203

it is a good habit to decorate function wrappers with @wraps;

instance method

we can have one instance method wrap another instance method in a different class like this:

class Foo:
    def foo(self, a0, a1=1, a2=2, a3=3):
        print(a0, a1, a2, a3)

class Bar:
    @wraps(Foo.foo)
    def bar(self, *args, **kwargs):
        sig = signature(Foo.foo)
        ba = sig.bind(self, *args, **kwargs)
        ba.arguments['a2'] = 102        #   set `a2` to `102`;
        Foo.foo(*ba.args, **ba.kwargs)

output:

>>> Foo().foo(200, a2=202, a3=203)
200 1 202 203
>>> Bar().bar(200, a2=202, a3=203)
200 1 102 203

if Bar extends Foo, the procedure can be further simplified:

class Foo:
    def foo(self, a0, a1=1, a2=2, a3=3):
        print(a0, a1, a2, a3)

class Bar(Foo):
    @wraps(Foo.foo)
    def bar(self, *args, **kwargs):
        sig = signature(self.foo)
        ba = sig.bind(*args, **kwargs)
        ba.arguments['a2'] = 102        #   set `a2` to `102`;
        self.foo(*ba.args, **ba.kwargs)

output:

>>> Foo().foo(200, a2=202, a3=203)
200 1 202 203
>>> Bar().bar(200, a2=202, a3=203)
200 1 102 203