python - 如何迭代一个既可以是整数又可以是数组的变量？

Question

我想制作一个同时接受整数和整数数组的多用途函数，例如一些 Numpy 函数，例如 delete做:

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

问题是当 bar 是单个 int 时，这显然会引发 TypeError: 'int' object is not iterable。而且我无法找到如何将 bar 转换为数组或任何可迭代的内容，并且在 bar 是数组时不会破坏代码。这个怎么做？

Answer 1

numpy.array有一个可选的 ndmin 参数。将其设置为 1 以保证您正在迭代一维数组:

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

您还可以指定 copy=False 以避免在输入已经是数组时生成副本。

请注意，我还修改了您的逻辑:您希望在迭代时打印 i 而不是 bar。