test_nditer.py 文件源码

def test_iter_allocate_output_subtype():
    # Make sure that the subtype with priority wins

    # matrix vs ndarray
    a = np.matrix([[1, 2], [3, 4]])
    b = np.arange(4).reshape(2, 2).T
    i = nditer([a, b, None], [],
                    [['readonly'], ['readonly'], ['writeonly', 'allocate']])
    assert_equal(type(a), type(i.operands[2]))
    assert_(type(b) != type(i.operands[2]))
    assert_equal(i.operands[2].shape, (2, 2))

    # matrix always wants things to be 2D
    b = np.arange(4).reshape(1, 2, 2)
    assert_raises(RuntimeError, nditer, [a, b, None], [],
                    [['readonly'], ['readonly'], ['writeonly', 'allocate']])
    # but if subtypes are disabled, the result can still work
    i = nditer([a, b, None], [],
            [['readonly'], ['readonly'], ['writeonly', 'allocate', 'no_subtype']])
    assert_equal(type(b), type(i.operands[2]))
    assert_(type(a) != type(i.operands[2]))
    assert_equal(i.operands[2].shape, (1, 2, 2))