How to rewrite this python function to accept lists as arguments?

Question

I'm fairly new to python and trying to figure out how to rewrite my squaring function here below to accept lists as arguments but I can't get it to work. I think I need to use things called map or reduce or something. Does anyone know how this should be rewritten to accept lists??

def square(self,x):
    number_types = (int, long, float, complex)
    if isinstance(x, number_types):
        return x*x
    else:
        raise ValueError

Answer 1

Just use a combination of lambda and map

 l=[1,2,3,4,5,6]


a= lambda x: x*x if type(x) == (int or float or complex) else ""

b=list(map(a,l))

print(b)

[1, 4, 9, 16, 25, 36]

Answer 2

To square each element:

def square(l):
   return [pow(i, 2) for i in l]

print(square([i for i in range(10)]))

Output:

[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Answer 3

Make a wrapper if you really need such function:

def square2(x):
    if isinstance(x, collections.abc.Sequence):
        return [square(n) for n in x]
    return square(x)

Change to collections.Sequence for Python < 3.3

Answer 4

Is this what you want:

class A:
    def square(self, x):
        number_types = (int, long, float, complex)
        xlist = x if isinstance(x, list) else [x]
        for i in xlist:
            if not isinstance(i, number_types):
                raise ValueError
        if not isinstance(x, list):
            return x*x
        return [i*i for i in x]

if __name__ == '__main__':
    a = A()
    print(a.square(2))
    print(a.square([1, 2, 3, 4, 5]))

Answer 5

Use NumPy

The best solution here would be to use numpy:

import numpy as np

def square(x):
    x = np.asarray(x)
    number_types = (int, long, float, complex)
    if x.dtype in number_types:
        return x*x
    else:
        raise ValueError

This is both faster than operating on lists, and allows you to work with any type of iterable. The modification to your code is also very small and the code is quite readable, especially when compared to map based solutions.

Examples

Works as expected with scalars:

>>> square(3)
9

Also works with lists, tuples, etc

>>> square([3, 4])
array([ 9, 16])
>>> square((3, 4))
array([ 9, 16])

Performance

A quick comparision of this to the other versions shows it is much faster

>>> a = lambda x: x*x if type(x) == (int or float or complex) else ""
>>> l = [0] * 100
>>> %timeit list(map(a,l))
10000 loops, best of 3: 23.5 µs per loop

>>> %timeit square(l)
100000 loops, best of 3: 6.88 µs per loop

For larger lists, the performance lead will be even larger.