How to iterate over a variable that can be both integer and array?

Question

I'd like to make a multi-purpose function accepting both integer, and array of integers, like how e.g. some Numpy function like delete does:

def foo(bar):   # type(foo) can be integer or an array of integers
    for i in bar:
        print(bar)

The problem is when bar is a single int, this obivously raises a TypeError: 'int' object is not iterable. And I couldn't find how to convert bar to an array, or anything iterable, in a way that doesn't break the code when bar is an array. How to do this?

Answer 1

numpy.array has an optional ndmin parameter. Set this to 1 to guarantee you are iterating a 1d array:

def foo(bar):   # type(bar) can be integer or an array of integers
    for i in np.array(bar, ndmin=1):
        print(i)

You can also specify copy=False to avoid a copy being produced if the input is already an array.

Note I have also amended your logic: you want to print i rather than bar when you iterate.