Decorator to allow function to accept arbitrary arguments

Question

How can I write a decorator so that decorated functions can accept (and ignore) arbitrary arguments?

I have some functions like this:

def foo(x):
    return x

def foo2(x, y):
    if bar(y):
        return x
    else:
        return x + 1

def foo3(x, y, z):
    ...

foo() can calculate the return value for a given x just based on x , but foo2() needs another parameter, and foo3() needs a third parameter. I have a method elsewhere that, among other things, calls either foo() , foo2() , etc depending on a user-specified argument.

Right now, the method just grabs the appropriate function with getattr(user_arg) and calls it with all of x , y , z . To avoid a TypeError for the wrong number of arguments, the foo functions are all defined with *args like this:

def foo(x, *args):
    return x

But I'd like to just have a decorator to avoid including *args in every single foo function definition. Is there a way to do that? Or can you suggest a better way to organize this code?

One way would be to write the method like this:

if user_arg == 'foo':
    foo(x)
elif user_arg == 'foo2':
    foo2(x, y)
elif user_arg == 'foo3':
    foo3(x, y, z)

But I'd really like to avoid this, because there are a lot of foo functions, and because I'd like to make it possible to add new foo functions without also having to add another branch to the method.

Edit : To clarify, it's not the case that I need to call a different foo function based on the number of arguments. It's arbitrary (user-specified) which of the foo functions is called.

def foo3(x, y, z):
    return x + y + z

def foo4(x, y, z):
    return x + y - z

Answer 1

You can use inspect.getargspec to determine what arguments the decorated function takes:

import functools
import inspect

def ignore_extra_arguments(func):
    args, varargs, kwvarargs, defaults = inspect.getargspec(func)
    @functools.wraps(func)
    def wrapper_func(*wrapper_args, **wrapper_kwargs):
        if varargs is None:
            # remove extra positional arguments
            wrapper_args = wrapper_args[:len(args)]
        if kwvarargs is None:
            # remove extra keyword arguments
            wrapper_kwargs = {k: v for k, v in wrapper_kwargs.iteritems() if k in args}
        return func(*wrapper_args, **wrapper_kwargs)
    return wrapper_func

This can be optimized a little by lifting the if checks out of wrapper_func , at the expense of writing more code.

Answer 2

You can use named arguments and always call foo() (the first function) in order to dispatch to the "right" function, as follows:

def foo(a, b=None, c=None, d=None):
    if b:
        return foo2(a, b, c, d)
    else:
        # do your thing

def foo2(a, b, c=None, d=None):        
    if c:
        return foo3(a, b, c, d)
    else:
        # do your thing

def foo3(a, b, c, d=None):                
    ...

---

Another option:

def foo(a, *b):
    if b:
        return foo2(a, *b)
    else:
        print "in foo"
        # do your thing

def foo2(a, b, *c):       
    if c:
        return foo3(a, b, *c)
    else:
        print "in foo2"
        # do your thing

def foo3(a, b, c, *d):
    print "in foo3"


foo(1, 2, 3)