CODEWITHSUNDEEP
pythonfunctionnumpyarguments
heltonbiker
heltonbiker

Reputation: 27605

Python function to accept numpy ndarray or sequence as arguments

I have seen some python functions that work generally receiving a (n,2) shaped numpy ndarray as argument, but can also "automagically" receive (2,n) or even len(2) sequences (tuple or list).

How is it pythonically achieved? Is there a unified good practice to check and treat these cases (example, functions in numpy and scipy module), or each developer implements which he thinks best?

I'd just like to avoid chains of (possibly nested) chains of ifs/elifs, in case there is a well known better way.

Thanks for any help.

Upvotes: 2

Views: 12265

Answers (1)

jterrace
jterrace

Reputation: 67083

You can use the numpy.asarray function to convert any sequence-like input to an array:

>>> import numpy
>>> numpy.asarray([1,2,3])
array([1, 2, 3])
>>> numpy.asarray(numpy.array([2,3]))
array([2, 3])
>>> numpy.asarray(1)
array(1)
>>> numpy.asarray((2,3))
array([2, 3])
>>> numpy.asarray({1:3,2:4})
array({1: 3, 2: 4}, dtype=object)

It's important to note that as the documentation says No copy is performed if the input is already an ndarray. This is really nice since you can pass an existing array in and it just returns the same array.

Once you convert it to a numpy array, just check the length if that's a requirement. Something like:

>>> def f(x):
...    x = numpy.asarray(x)
...    if len(x) != 2:
...       raise Exception("invalid argument")
... 
>>> f([1,2])
>>> f([1,2,3])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 4, in f
Exception: invalid argument

Update: Since you asked, here's a "magic" function that will except *args as an array also:

>>> def f(*args):
...    args = numpy.asarray(args[0]) if len(args) == 1 else numpy.asarray(args)
...    return args
... 
>>> f(7,3,5)
array([7, 3, 5])
>>> f([1,2,3])
array([1, 2, 3])
>>> f((2,3,4))
array([2, 3, 4])
>>> f(numpy.array([1,2,3]))
array([1, 2, 3])

Upvotes: 6

Related Questions