Python function that handles scalar or arrays

All over the numpy code base you find things like:

def func_for_scalars_or_vectors(x):    x = np.asarray(x)    scalar_input = False    if x.ndim == 0:        x = x[None]  # Makes x 1D        scalar_input = True    # The magic happens here    if scalar_input:        return np.squeeze(ret)    return ret

As a matter of opinion, I would prefer to have a function be flexible on input types, but always return a consistent type; this is what will ultimately prevent callers from having to check return types (the stated goal).

For example, allow scalars and/or arrays as arguments, but always return the array.

def func(x, y):    # allow (x=1,y=2) OR (x=[1,2], y=[3,4]) OR (!) (x=1,y=[2,3])    xn = np.asarray([x]) if np.isscalar(x) else np.asarray(x)    yn = np.asarray([y]) if np.isscalar(y) else np.asarray(y)    # calculations with numpy arrays xn and xy    res = xn + yn  # ..etc...    return res

(Still, the example can easily be modified to return a scalar, by setting a flag "scalar=True", yada yada yada.. but you'd also have to handle one arg's a scalar, the other is an array, etc.; seems like a lot of YAGNI to me.)

" function that can accept either scalar floats or numpy vectors (1-d array), and return a scalar, 1-d array, or 2-d array"

So

scalar => scalar

1d => 2d

what produces a 1-d array?

def func( xa, ya ):    def something_complicated(x, y):        return x + y    try:        xy = np.zeros( ( len(xa), len(ya) ) )        for j in range(len( ya )):            for i in range(len( xa )):                xy[i,j] = something_complicated(xa[i], ya[i])    except TypeError:        xy = something_complicated(xa, ya)      return xy

Is this ' fast, elegant, and pythonic'?

It certainly is 'pythonic'. 'try/except' is very Pythonic. So is defining a function within another function.

Fast? Only time tests will tell. It may depend on the relative frequency of scalar v. array examples.

Elegant? That is in the eyes of the beholder.

Is this more elegant? It's limited recursion

def func( xa, ya ):    try:        shape = len(xa), len(ya)    except TypeError:        # do something complicated        return xa+ya        xy = np.zeros(shape)    for j in range(len( ya )):        for i in range(len( xa )):            xy[i,j] = func(xa[i], ya[i])               return xy

If you need to correctly handle 2d+ inputs, then vectorize is clearly the least effort solution:

def something_complicated(x,y):    return x+yvsomething=np.vectorize(something_complicated)In [409]: vsomething([1,2],[4,4])Out[409]: array([5, 6])In [410]: vsomething(1,3)Out[410]: array(4)   # not quite a scalar

If array(4) is not the scalar output that you want, then you'll have to add a test and extract the value with [()]. vectorize also handles a mix of scalar and array (scalar + 1d => 1d).

MATLAB does not have scalars. size(3) returns 1,1.

In Javascript, [1,2,3] has a .length attribute, but 3 does not.

from a nodejs session:

> x.lengthundefined> x=[1,2,3][ 1, 2, 3 ]> x.length3

Regarding MATAB code, Octave has this to say about the length function

-- Built-in Function: length (A) Return the length of the object A.

The length is 0 for empty objects, 1 for scalars, and the number of elements for vectors. For matrix objects, the length is the number of rows or columns, whichever is greater (this odd definition is used for compatibility with MATLAB).

MATLAB does not have true scalars. Everything is at least 2d. A 'vector' just has a '1' dimension. length is a poor choice for iteration control in MATLAB. I've always used size.

To add to the MATLAB convenience, but also potential confusion, x(i) works with both row 'vectors', and column 'vectors', [1,2,3] and [1;2;3]. x(i,j) also works with both, but with different index ranges.

len works fine when iterating for Python lists, but isn't the best choice when working with numpy arrays. x.size is better if you want to total number of items. x.shape[0] is better if you want the 1st dimension.

Part of why there isn't an elegant Pythonic solution to your problem is that you are starting with some that is idiomatic MATLAB, and expected Python to behave with all the same nuances.