overloading a float to a numpy array

问题

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?

回答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])

来源：https://stackoverflow.com/questions/18785237/overloading-a-float-to-a-numpy-array