test_nditer.py 文件源码

def test_0d_nested_iter():
    a = np.arange(12).reshape(2, 3, 2)
    i, j = np.nested_iters(a, [[], [1, 0, 2]])
    vals = []
    for x in i:
        vals.append([y for y in j])
    assert_equal(vals, [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])

    i, j = np.nested_iters(a, [[1, 0, 2], []])
    vals = []
    for x in i:
        vals.append([y for y in j])
    assert_equal(vals, [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11]])

    i, j, k = np.nested_iters(a, [[2, 0], [], [1]])
    vals = []
    for x in i:
        for y in j:
            vals.append([z for z in k])
    assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])