overloading a float to a numpy array

Question

I have a function, processing a 1D numpy array, like this:

def f(arr):
 arr=asarray(arr)
 #process data as numpy array
 #...
 return arr

With asarray I allow to call the function with a list as f([4,5,6]). Now, I would like to "overload" the argument also to a single float, so that I can use f(4) instead of f([4]).

This is a standard numpy feature, since you can call np.sin as sin(array([4,5,6])), or as sin([4,5,6]) or as sin(4) as well. I came up with this code, that works at least in simple cases:

def f(arr):
 arr=asarray(arr)
 if arr.shape is ():
 print 'arr is a single float/int/etc'
 arr = array([arr])
 #process data as numpy array
 #...
 return arr

Is this the standard/correct way to do it?

Answer 1

I believe you are looking for np.atleast_1d.

>>> np.atleast_1d(5)
array([5])
>>> np.atleast_1d(np.arange(2))
array([0, 1])