[Numpy-discussion] array_arguments decorator

Tim Hochberg tim.hochberg at cox.net
Fri Mar 31 19:49:03 CST 2006

Zachary Pincus wrote:

> Hi folks -
> I had seen some talk on this list about the utility of a decorator  
> for functions that need to convert their arguments to numpy arrays.  
> This would help eliminate boilerplate calls to 'asarray' like:
> def distance_squared(a, b):
>   a = numpy.asarray(a)
>   b = numpy.asarray(b)
>   return ((a - b)**2).sum()
> Here is a trivial decorator I was thinking of adding to the wiki --  
> does this cover enough cases to be useful? In a bigger sense, would  
> it be worthwhile to add some decorators like this to numpy itself?  
> (I'm not sure I'm in favor of this, since I kind of like smaller APIs  
> over bigger ones.)
> def array_arguments(f):
>   """Wrap a function such that any positional arguments are
>   converted into numpy arrays before the function is called."""
>   def convert_arg_wrapper(*args, **kwargs):
>     array_args = [numpy.asarray(a) for a in args]
>     return f(*array_args, **kwargs)
>   return convert_arg_wrapper
> now distance_squared can look like:
> @array_arguments
> def distance_squared(a, b):
>   return ((a - b)**2).sum()
> if using python 2.4, or if not so using:
> def distance_squared(a, b):
>   return ((a - b)**2).sum()
> distance_squared = array_arguments(distance_squared)

Great minds think alike. Or at least our minds think alike ;) I also 
wrote up a decorator for this same purpose. Then I got distracted and 
forgot to post it. Mine has more features at the expense of being more 
complicated. The main extra feature is that it allows you to decide both 
which args get checked and what there types should be. It also preserves 
the signature of the original function. Some of this stuff is 
accomplished using the decorator decorator, which you can find here:


The upshot of all of this is that you can do stuff like:

        @array_function(a=float, b=None, c=complex)
        def foo(a, b, c=1, d=None):
            print repr(a), repr(b), repr(c), repr(d)

And it will convert 'a' and 'c' to float and complex arrays respectively 
and convert 'b' to some type of array. Arguments not mentioned don't get 
touched, unless you specify no arguments, in which case all of the 
positional arguments get converted (*args and **kwargs are not touched 
in this case).

I'm not certain that this is the world's best interface, nor am I 
certain that the extra complexity is worth it -- yours is certainly 
easier to understand. However, it was fun to write. I am unlikely to 
find the time to do anything with it anytime soon, so should you desire 
to revamp and place it on the wiki, feel free. Or, if you want to ignore 
it, feel free to do that as well.

The code is below:



import inspect, decorator, numpy

def array_function(**signature):
    def deco(func):
        regargs, varargs, varkwargs, defaults = inspect.getargspec(func)
        if not signature:
            signature.update((name, None) for x in regargs)
        def caller(func, *args, **kwargs):
            args = list(args)
            for i, (name, value) in enumerate(zip(regargs, args)):
                if name in signature:
                    args[i] = numpy.asarray(value, signature[name])
            for name, value in kwargs.items():
                if name in signature:
                    kwargs[name] = numpy.asarray(value, signature[name])
            return func(*args, **kwargs)
        return decorator._decorate(func, caller)
    return deco

More information about the Numpy-discussion mailing list