1

In Python, I would like to have a function working on different input types. Something like this:

def my_square(x):
    return x ** 2

my_square(2)             #return 4
my_square(range(10))     #should return a list [0 ... 81]

npa = numpy.zeros(10)
my_square(npa)           # should return a numpy array with the squares of zeros

Basically, what is good practice to write functions for both scalars and iterables? Can this be done with *args or *kwargs perhaps?

RickB
  • 175
  • 1
  • 11
  • 7
    Why not write the function to work only on scalars then e.g. `map` it onto iterables? You certainly *can* handle both within the function, but I'd try to avoid it. – jonrsharpe Jun 26 '15 at 15:12
  • It's a bad idea to do what you're trying to do, except in fringe cases. Write versions for each type of input. – Carcigenicate Jun 26 '15 at 15:15
  • @jonrsharpe Numpy arrays are designed so that expressions can be written that work whether the variables used within are scalars or arrays. The OP's `my_square(x)` will work equally as well whether `x = 4` or `x` is an `ndarray` of one million entries. In this case, writing a function that accepts either a scalar or an iterable is a common numerical python idiom used in many scientific packages, and using `map` to apply it to an iterable is decidedly an anti-pattern. But if the function is general (and can't take advantage of numpy's vectorization), then a list comp. or map is the way to go. – jme Jun 26 '15 at 16:12
  • @jme interesting, thanks! – jonrsharpe Jun 26 '15 at 16:14

4 Answers4

0

A typical way to do this is to use numpy.asarray to convert the input of your function to an ndarray. If the input is a scalar, then the result is an array with 0 dimensions, which acts essentially like a Python number type. For instance:

def square(x):
    x = np.asarray(x)
    return x**2

So that:

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

Note that I gave a list as input and received an ndarray as output. If you absolutely must receive the same type of output as you provided as input, you can convert the result before returning it:

def square(x):
    in_type = type(x)
    x = np.asarray(x)
    return in_type(x**2)

But this incurs an additional cost for little benefit.

jme
  • 19,895
  • 6
  • 41
  • 39
0

Since Python is dynamically typed, and the design philosophy of Python is against the idea of wanting a difference to be obvious when the function is called, this isn't considered Pythonic. You can however use the isInstance() method to accomplish what you want, if you must that is. For example:

def mySquare(x):
    if isinstance(x, int):
        return x**2
    elif isinstance(x, range):
        return [i ** 2 for i in x]

print(mySquare(2)); //4
print(mySquare(range(10))); //[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

IDEOne link here.

However, just because you can do something, doesn't mean you should.

Refer to this question for further information on isinstance, and I suggest you take a look at Duck Typing as well.

Additionally, a single dispatch function might also provide what you need, but I am not experienced enough to provide an explanation for this, however, it might be something you want to look into.

Community
  • 1
  • 1
matrixanomaly
  • 6,627
  • 2
  • 35
  • 58
  • This `isinstance` approach is very fragile - what if the input is a `float`? It would be better to use e.g. `collections.abc.Sequence`, or duck typing (`try` to iterate over the input, and assume it's scalar if that fails). – jonrsharpe Jun 26 '15 at 15:32
  • @jonrsharpe I agree, I'm trying to make it work for numpy.zeros and I can't get a right result given how isinstance conditionals will not execute if they don't match exactly. definitely unpythonic, but quick and dirty. – matrixanomaly Jun 26 '15 at 15:38
  • 1
    You can provide multiple classes, e.g. `isinstance(x, (int, float))`, but that can end up being more trouble than it's worth. – jonrsharpe Jun 26 '15 at 15:39
0

It would be much better practice and much more maintainable to just do:

def foo(n):
    return n ** 2

And then build lists, dicts, etc when needed (I'm not super familiar with numpy but I imagine there is something similar you could do):

foo_list = [foo(n) for n in range(10)]
foo_dict = {n: foo(n) for n in range(10)}

And it seems that there is for numpy using numpy.fromiter(). From the docs:

iterable = (foo(n) for n in range(5))
foo_arr = np.fromiter(iterable, np.whatever_numpy_datatype)

you could even make those into functions if you really needed to:

def foo_list(start=0, stop=0):
    return [foo(n) for n in range(start, stop)]

def foo_dict(start=0, stop=0):
    return {n: foo(n) for n in range(10)}

Another option as it's easier to ask for forgiveness than permission:

def foo(scalar_or_iter):
    try:
        return [n ** 2 for n in scalar_or_iter]
    except TypeError:
        return scalar_or_iter ** 2
kylieCatt
  • 10,672
  • 5
  • 43
  • 51
  • To get the square of the elements of a NumPy array `x`, you just write `x**2`. It would be horribly inefficient to square a NumPy array via a Python-level iteration. – Mark Dickinson Jun 27 '15 at 15:11
  • So it would be more efficient to just make an array like so `foo_arr = np.asarray(some_list)` and then just `foo_arr**2` @MarkDickinson? – kylieCatt Jun 27 '15 at 15:20
  • Perhaps. If you're really starting with a list, you'll need to decide whether it's really worth converting to an array or not. My point was that if you're *starting* with a NumPy array, it doesn't make sense to iterate over its elements at Python level (and in that case there's no need for the `asarray` call). – Mark Dickinson Jun 27 '15 at 16:01
0

As suggested in the comments, just write a scalar version of your function and use map for lists etc and imap for iterables (map will not work on those):

map(myFunc,myList)

and

import itertools
itertools.imap(myFunc,myIterable)
Stijn Nevens
  • 246
  • 2
  • 5