How to make a function that can handle single inputs or lists of inputs

Question

I'll illustrate my question with a couple of overly simple functions that just square numbers. But it is the general best practice for writing functions that I want to know about.

Suppose I have a number that I want to square, I could write the following function.

def square (x):
    y = x**2
    return y

If I had a list of numbers and wanted to get the list in which each element is squared, I could use

def square (x_list):
    y_list = []
    for x in x_list:
        y_list.append( x**2 )
    return y_list

But I want to write just one function that will handle both of these cases. It will see whether the input is a number or a list, and act accordingly. I could do this using type to detect type, but I'd like to know what the most pythonic way to do it is.

Answer 1

You can check for the type of the passed in argument and act accordingly:

# check the variable instance type
def square (x):
    """Returns squares of ints and lists of ints - even if boxed inside each other.
    It uses recursion, so do not go too deep ;o)"""

    if isinstance(x,int):
        return x**2

    elif isinstance(x,list):
        return [square(b) for b in x] #  recursion for lists of ints/lists of ints

    else:
        raise ValueError("Only ints and list of ints/list of ints allowed")


print(square(4))
print(square([2,3,4]))
print(square([2,3,[9,10],11]))

try:
    print(square(2.6))
except ValueError as e:
    print(e)

Output:

16
[4, 9, 16]
[4, 9, [81, 100], 121]

Only ints and list of ints/list of ints allowed

Answer 2

I would agree with @Bryan Oakley's answer, that it's best to write the code to only accept a single argument type. With that being said, I figured I would present an example of a function that handles a variable number of input arguments:

def square(*args):

    return [arg**2 for arg in args]

Note that this will always return a list:

y = square(2,4,6,8)

y = square(4)

Yields:

[4, 16, 36, 64]
[16]