编写一个可以接受值和任何迭代类型的函数

时间:2017-05-05 13:16:49

标签: python list numpy iterator

我想知道是否可以用非常pythonic的方式编写以下代码。 代码被简化了。它应该接受每个输入的单个值,以及序列或numpy数组。问题是单个值没有任何len()参数。目前我正在检查类型是numpy array还是list

但是将来,可能会有不同的数据表示将被迭代,如熊猫阵列或类似的。这将导致代码的更改和更长的if语句。

如果可以检查输入是否可迭代,那将是非常好的,以便代码将来是动态的。

简化示例代码:

import numpy as np


def test(a,b):
    c = []
    if type(a) != type(b):
        return None;

    if type(a) is np.ndarray or type(a) is list:
        for i in range(0,len(a)):
            c.append(a[i] + b[i])
    else:
        c = a + b

    return c

A = test(1,2)
B = test([1,2,3],[1,2,3])
C = test(np.asarray([1,2,3]),np.asarray([1,2,3]))

感谢您的帮助

1 个答案:

答案 0 :(得分:1)

您可以检查两个元素是否都有__iter__方法。

def test(a, b):
    if not hasattr(a, '__iter__') and not hasattr(b, '__iter__'):
        return a + b
    return [x + y for x, y in zip(a, b)]  # Let zip() raise the TypeError