test_nditer.py 文件源码

def test_iter_nested_iters_dtype_copy():
    # Test nested iteration with a copy to change dtype

    # copy
    a = arange(6, dtype='i4').reshape(2, 3)
    i, j = np.nested_iters(a, [[0], [1]],
                        op_flags=['readonly', 'copy'],
                        op_dtypes='f8')
    assert_equal(j[0].dtype, np.dtype('f8'))
    vals = []
    for x in i:
        vals.append([y for y in j])
    assert_equal(vals, [[0, 1, 2], [3, 4, 5]])
    vals = None

    # updateifcopy
    a = arange(6, dtype='f4').reshape(2, 3)
    i, j = np.nested_iters(a, [[0], [1]],
                        op_flags=['readwrite', 'updateifcopy'],
                        casting='same_kind',
                        op_dtypes='f8')
    assert_equal(j[0].dtype, np.dtype('f8'))
    for x in i:
        for y in j:
            y[...] += 1
    assert_equal(a, [[0, 1, 2], [3, 4, 5]])
    i, j, x, y = (None,)*4  # force the updateifcopy
    assert_equal(a, [[1, 2, 3], [4, 5, 6]])