Python Hackery: merging signatures of two Python functions

Starting point

def old(foo, bar):
    """This is old's docstring."""
    print(foo, bar)
    return foo + bar


def new(prefix, foo, *args, **kwargs):
    return old(prefix + foo, *args, **kwargs)

Let’s test it.

>>> o = old('a', 'b')
a b
>>> n = new('!', 'a', 'b')
!a b
>>> print(o, n, sep=' - ')
ab - !ab
>>> help(old)
Help on function old in module __main__:

old(foo, bar)
    This is old's docstring.

>>> help(new)
Help on function new in module __main__:

new(prefix, foo, *args, **kwargs)

The last line is not exactly informative — it doesn’t tell us that we need to pass bar as an argument. Sure, you could define new as just (prefix, foo, bar) — but that means every change to old requires editing new as well. So, not ideal. Let’s try to fix this.

The existing infrastructure: functools.wraps

First, let’s start with the basic facility Python already has. The standard library already comes with functools.wraps and functools.update_wrapper.

If you’ve never heard of those two functions, here’s a crash course:

def decorator(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        print("Inside wrapper")
        f(*args, **kwargs)
    return wrapper

@decorator
def square(n: float) -> float:
    """Square a number."""
    return n * n

If we try to inspect the square function, we’ll see the original name, arguments, annotations, and the docstring. If we ran this code again, but with the @functools.wraps(f) line commented out, we would only see wrapper(*args, **kwargs).

This approach gives us a hint of what we need to do. However, if we apply wraps (or update_wrapper, which is what wraps ends up calling) to our function, it will only have foo and bar as arguments, and its name will be displayed as old.

So, let’s take a look at functools.update_wrapper. What does it do? Two things:

• copy some attributes from the old function to the new one (__module__, __name__, __qualname__, __doc__, __annotations__)

• update __dict__ of the new function

• set wrapper.__wrapped__

If we try to experiment with it — by changing the list of things to copy, for example — we’ll find out that the annotations, the docstring, and the displayed name come from the copied attributes, but the signature itself is apparently taken from __wrapped__.

Further investigation reveals this fact about inspect.signature:

inspect.signature(callable, *, follow_wrapped=True)

New in version 3.5: follow_wrapped parameter. Pass False to get a signature of callable specifically (callable.__wrapped__ will not be used to unwrap decorated callables.)

And so, this is our end goal:

Craft a function with a specific signature (that merges old and new) and set it as new.__wrapped__.

But first, we need to talk about parallel universes.

Or actually, code objects.

Defining a function programmatically

Let’s try an experiment.

>>> def foo(bar): pass
>>> foo.__wrapped__ = lambda x, y: None
>>> help(foo)
foo(x, y)

So, there are two ways to do this. The first one would be to generate a string with the signature and just use eval to get a __wrapped__ function. But that would be cheating, and honestly, quite boring. (The inspect module could help us with preparing the string.) The second one? Create code objects manually.

code(argcount, kwonlyargcount, nlocals, stacksize, flags, codestring,
    constants, names, varnames, filename, name, firstlineno,
    lnotab[, freevars[, cellvars]])

Create a code object.  Not for the faint of heart.

>>> def f(a, b=1, c=2, *, d=3): pass
>>> inspect.getfullargspec(f)
FullArgSpec(args=['a', 'b', 'c'], varargs=None, varkw=None, defaults=(1, 2), kwonlyargs=['d'], kwonlydefaults={'d': 3}, annotations={})

Final results

Before I show you the code, let’s test it out:

# old defined as before
@merge_args(old)
def new(prefix, foo, *args, **kwargs):
    return old(prefix + foo, *args, **kwargs)

And the end result — help(new) says:

new(prefix, foo, bar)
    This is old's docstring.

We did it!

The code is available on GitHub and on PyPI (pip install merge_args). There’s also an extensive test suite.

PS. you might be interested in another related post of mine, in which I reverse-engineer the compilation of a function: Gynvael’s Mission 11 (en): Python bytecode reverse-engineering