Python函数接受numpy ndarray或sequence作为参数

Question

我见过一些python函数，它们通常接收（n，2）形状的numpy ndarray作为参数，但也可以“自动”接收（2，n）甚至len（2）序列（元组或列表）。

它是怎样达到的呢？是否有一个统一的良好实践来检查和处理这些情况（例如，numpy和scipy模块中的函数），或者每个开发人员都实现了他认为最好的功能？

我只想避免ifs/elif的链（可能是嵌套的），以防有一种众所周知的更好的方法。

谢谢你的帮助。

Answer 1

可以使用numpy.asarray函数将任何类似序列的输入转换为数组：

>>> import numpy
>>> numpy.asarray([1,2,3])
array([1, 2, 3])
>>> numpy.asarray(numpy.array([2,3]))
array([2, 3])
>>> numpy.asarray(1)
array(1)
>>> numpy.asarray((2,3))
array([2, 3])
>>> numpy.asarray({1:3,2:4})
array({1: 3, 2: 4}, dtype=object)

需要注意的是，正如文档中所说，如果输入已经是ndarray，则不执行任何拷贝。这真的很好，因为您可以传入一个现有数组，它只返回相同的数组。

一旦将其转换为numpy数组，只要检查长度（如果需要）。类似于：

>>> def f(x):
...    x = numpy.asarray(x)
...    if len(x) != 2:
...       raise Exception("invalid argument")
... 
>>> f([1,2])
>>> f([1,2,3])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 4, in f
Exception: invalid argument

更新： 既然你问了，这里有一个“magic”函数，除了*args作为数组之外：

>>> def f(*args):
...    args = numpy.asarray(args[0]) if len(args) == 1 else numpy.asarray(args)
...    return args
... 
>>> f(7,3,5)
array([7, 3, 5])
>>> f([1,2,3])
array([1, 2, 3])
>>> f((2,3,4))
array([2, 3, 4])
>>> f(numpy.array([1,2,3]))
array([1, 2, 3])