[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:
http://www.phyast.pitt.edu/~micheles/python/
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:
Regards,
-tim
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