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

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]

Answer 2

It's very simple if you use numpy. For multiple elements make it as an array. For single element ,the function is applied as it.

import numpy as np

def square(a):
    return a**2

a = np.array([2,3,4,5])
print(square(a)) # array([ 4,  9, 16, 25])

b = 9
print(square(b)) # 81

Answer 3

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 4

Treat the argument as list input and handle the exception if it's not.

EAFP Easier to ask for forgiveness than permission. This common Python coding style assumes the existence of valid keys or attributes and catches exceptions if the assumption proves false. This clean and fast style is characterized by the presence of many try and except statements. The technique contrasts with the LBYL style common to many other languages such as C. docs

def square(x):    
    try:
        y = [e ** 2 for e in x]
    except TypeError:
        y = x ** 2
    return y

square(2)
# Out: 4
square([2, 3])
# Out: [4, 9]

Answer 5

As other answers have explained, this probably isn't a great design.

First, a "number" could be an int, or some user-defined subclass of int, or a float, or some user-defined Quaternion type. Normally, you just use duck-typing: it x ** 2 works, then x is quacking like a number, and that's good enough.

But a list of ints doesn't quack like an int. So, what can you do?

Well, usually, you'll want to explicitly loop over them:

>>> xs = [1, 2, 3]
>>> sqs = [square(x) for x in xs]

… or write a function that does that:

>>> def squares(xs): return [square(x) for x in xs]
>>> sqs = squares(xs)

… or use a type that knows how to vectorize mathematical operators:

>>> xs = np.array([1, 2, 3])
>>> sqs = square(xs)

---

In fact, even when you do want to handle two different types, you can often rely on duck typing:

def square(x):
    try:
        return x**2
    except ValueError:
        return [i**2 for i in x]

This will square anything that's number-like enough to be squarable, and iterate over squaring all of the elements of anything that isn't, and raise a reasonable exception for anything that fails (either because it's not iterable, or because its elements aren't squarable).

---

Occasionally, you really do need to type-switch. But you still want to keep as close to duck-typing as possible, and that means using isinstance (so that, e.g., a user subtype of int still counts as a number) and, usually, using abstract base classes (so that, e.g., a user Quaternion type still counts as a number).

In this case, that means either treating numbers specially and assuming anything else is an iterable:

def square(x):
    if isinstance(x, numbers.Number):
        return x**2
    else:
        return [i**2 for i in x]

… or treating iterables specially and assuming everything else is a number:

def square(x):
    if isinstance(x, collections.abc.Iterable):
        return [i**2 for i in x]
    else:
        return x**2

Or maybe treating both specially and calling everything else an error:

def square(x):
    if isinstance(x, numbers.Number):
        return x**2
    elif isinstance(x, collections.abc.Iterable):
        return [i**2 for i in x]
    raise ValueError(f"'{type(x).__name__}' instance is not a number or numbers")

Answer 6

You can write one line return but using type:

def square(x):
    return(x**2 if type(x) is int else [i**2 for i in x])