Clean way to return a one-element Numpy array

Question

Is there a clean way to write functions that return a one-element numpy array as the element itself?

Let's say I want to vectorize a simple square function and I want my return value to be the same dtype as my input. I could write something like this:

def foo(x):
    result = np.square(x)
    if len(result) == 1:
        return result[0]
    return result

or

def foo(x):
    if len(x) == 1:
        return x**2
    return np.square(x)

Is there an easier way to do this? So that I can use this function both for scalars and for arrays?

I know that I can check the dtype of my input directly and use IF statements to make it work, but is there a cleaner way?

Answer 1

I am not really sure whether or not I fully understood the question, but maybe something like this would help?

def square(x):
    if 'numpy' in str(type(x)):
        return np.square(x)
    else:
        if isinstance(x, list):
            return list(np.square(x))
        if isinstance(x, int):
            return int(np.square(x))
        if isinstance(x, float):
            return float(np.square(x))

I defined some test cases:

np_array_one = np.array([3.4])
np_array_mult = np.array([3.4, 2, 6])
int_ = 5
list_int = [2, 4, 2.9]
float_ = float(5.3)
list_float = [float(4.5), float(9.1), float(7.5)]

examples = [np_array_one, np_array_mult, int_, list_int, float_, list_float]

So we can see how the function behaves.

for case in examples:
    print 'Input type: {}.'.format(type(case))
    out = square(case)
    print out
    print 'Output type: {}'.format(type(out))
    print '-----------------'

And the output:

Input type: <type 'numpy.ndarray'>.
[ 11.56]
Output type: <type 'numpy.ndarray'>
-----------------
Input type: <type 'numpy.ndarray'>.
[ 11.56   4.    36.  ]
Output type: <type 'numpy.ndarray'>
-----------------
Input type: <type 'int'>.
25
Output type: <type 'int'>
-----------------
Input type: <type 'list'>.
[4.0, 16.0, 8.4100000000000001]
Output type: <type 'list'>
-----------------
Input type: <type 'float'>.
28.09
Output type: <type 'float'>
-----------------
Input type: <type 'list'>.
[20.25, 82.809999999999988, 56.25]
Output type: <type 'list'>
-----------------

From the test cases, the input and output is always same. However, the function is not really clean.

I used some of the code from this question at SO.