CODEWITHSUNDEEP
pythonnumpynumpy-ndarray
gkeenley
gkeenley

Reputation: 7418

How to handle both float and array input in Python?

I'm trying to write a function that accepts either a float OR an array of floats, and handles both of them using the same lines of code. For example, I want to return the float itself if it's a float, and the sum of the array of floats if it's an array. Something like this

def func(a):
  return np.sum(a)

and have both func(1.2) return 1.2, and func(np.array([1.2,1.3,1.4]) return 3.9.

Upvotes: 0

Views: 1465

Answers (4)

Nils Werner
Nils Werner

Reputation: 36849

The usual way to make sure the input is a NumPy array is to use np.asarray():

import numpy as np

def func(a):
  a = np.asarray(a)
  return np.sum(a)

func(1.2)
# 1.2
func([1.2, 3.4])
# 4.6
func(np.array([1.2, 3.4]))
# 4.6

or, if you want to get len() of your array, make sure it is at least 1-dimensional, use np.atleast_1d():

def func(a):
  a = np.atleast_1d(a)
  return a.shape[0]

func(1.2)
# 1
func([1.2, 3.4])
# 2
func(np.array([1.2, 3.4]))
# 2

Upvotes: 2

Sam Chats
Sam Chats

Reputation: 2311

You can use argument flattening:

def func(*args):
    # code to handle args
    return sum(args)

Now the following have the same behaviour:

>>> func(3)
3
>>> func(3, 4, 5)
12
>>> func(*[3, 4, 5])
12

Upvotes: 1

Rob Streeting
Rob Streeting

Reputation: 1735

You can check if the input is a float, and then put it in a list before processing the sum:

def func(a):
    if isinstance(a, float):
        a = [a]
    return np.sum(a)

Upvotes: 0

Celius Stingher
Celius Stingher

Reputation: 18377

This already works, where's the problem?

import numpy as np
def func(a):
  return np.sum(a)
print(func(np.array([1.2,2.3,3.2])))
print(func(1.2))

Output:

6.7
1.2

Upvotes: 1

Related Questions