decorating functions with generic signatures (not for the faint of heart)

[Michele Simionato]

I have realized today that defining decorators for functions
with generic signatures is pretty non-trivial.

Consider for instance this typical code:

#<traced_function.py>

def traced_function(f):
    def newf(*args, **kw):
        print "calling %s with args %s, %s" % (f.__name__, args, kw)
        return f(*args, **kw)
    newf.__name__ = f.__name__
    return newf

@traced_function
def f1(x):
    pass

@traced_function
def f2(x, y):
    pass

#</traced_function.py>

This is simple and works:

>>> from traced_function import traced_function, f1, f2
>>> f1(1)
calling f1 with args (1,), {}
>>> f2(1,2)
calling f2 with args (1, 2), {}

However, there is a serious disadvantage: the decorator replaces
a function with a given signature with a function with a generic
signature. This means that the decorator is *breaking pydoc*!

$ pydoc2.4 traced_function.f1
Help on function f1 in traced_function:

traced_function.f1 = f1(*args, **kw)

You see that the original signature of f1 is lost: even if I will get
an error when I will try to call it with a wrong number of arguments,
pydoc will not tell me that :-(

The same is true for f2:

$ pydoc2.4 traced_function.f2
Help on function f2 in traced_function:

traced_function.f2 = f2(*args, **kw)

In general all functions decorated by 'traced_function' will have the
same (too much) generic signature. This is a disaster for people
like me that rely heavily on Python introspection features.

I have found a workaround, by means of a helper function that
simplifies
the creation of decorators. Let's call this function 'decorate'.
I will give the implementation later, let me show how it works first.

'decorate' expects as input two functions: the first is the function
to be decorated (say 'func'); the second is a caller function
with signature 'caller(func, *args, **kw)'.
The caller will call 'func' with argument 'args' and 'kw'.
'decorate' will return a function *with the same signature* of
the original function, but enhanced by the capabilities provided
by the caller.

In our case we may name the caller function 'tracer', since
it just traces calls to the original function. The code makes
for a better explanation:

#<traced_function2.py>

from decorate import decorate

def tracer(f, *args, **kw):
    print "calling %s with args %s, %s" % (f.func_name, args, kw)
    return f(*args, **kw)

def traced_function(f):
    "This decorator returns a function decorated with tracer."
    return decorate(f, tracer)

@traced_function
def f1(x):
    pass

@traced_function
def f2(x, y):
    pass

#</traced_function2.py>

Let me show that the code is working:

>>> from traced_function2 import traced_function, f1, f2
>>> f1(1)
calling f1 with args (1,), {}
>>> f2(1,2)
calling f2 with args (1, 2), {}

Also, pydoc gives the right output:

$ pydoc2.4 traced_function2.f2
Help on function f1 in traced_function2:

traced_function2.f1 = f1(x)

$ pydoc2.4 traced_function2.f2
Help on function f2 in traced_function2:

traced_function2.f2 = f2(x, y)

In general all introspection tools using inspect.getargspec will
give the right signatures (modulo bugs in my implementation of
decorate).

All the magic is performed by 'decorate'. The implementation of
'decorate' is not for the faint of heart and ultimately it resorts
to 'eval' to generate the decorated function. I guess bytecode
wizards here can find a smarter way to generate the decorated function.
But my point is not about the implementation (which is very little
tested
at the moment). My point is that I would like to see something like
'decorate' in the standard library.
I think somebody already suggested a 'decorator' module containing
facilities to simplify the usage of decorators. This post is meant as
a candidate for that module. In any case, I think 'decorate' makes a
good example of decorator pattern.

Here is my the current implementation (not very tested):

#<decorate.py>

def _signature_gen(varnames, default_args, n_args, rm_defaults=False):
    n_non_default_args = n_args - len(default_args)
    non_default_names = varnames[:n_non_default_args]
    default_names = varnames[n_non_default_args:n_args]
    other_names = varnames[n_args:]
    n_other_names = len(other_names)
    for name in non_default_names:
        yield "%s" % name
    for name, default in zip(default_names, default_args):
        if rm_defaults:
            yield name
        else:
            yield "%s = %s" % (name, default)
    if n_other_names == 1:
        yield "*%s" % other_names[0]
    elif n_other_names == 2:
        yield "*%s" % other_names[0]
        yield "**%s" % other_names[1]

def decorate(func, caller):
    argdefs = func.func_defaults or ()
    argcount = func.func_code.co_argcount
    varnames = func.func_code.co_varnames
    signature = ", ".join(_signature_gen(varnames, argdefs, argcount))
    variables = ", ".join(_signature_gen(varnames, argdefs, argcount,
                                        rm_defaults=True))
    lambda_src = "lambda %s: call(func, %s)" % (signature, variables)
    dec_func = eval(lambda_src, dict(func=func, call=caller))
    dec_func.__name__ = func.__name__
    dec_func.__doc__ = func.__doc__
    dec_func.__dict__ = func.__dict__.copy()
    return dec_func

#</decorate.py>