def generic(self, args, kws):
assert not kws
if len(args) != 1:
return
arrays, = args
if isinstance(arrays, types.BaseTuple):
if not arrays:
return
arrays = list(arrays)
else:
arrays = [arrays]
nditerty = types.NumpyNdIterType(arrays)
return signature(nditerty, *args)
def test_properties(self):
def check(ty, dtypes, ndim, layout, indexers=None):
self.assertEqual(ty.ndim, ndim)
self.assertEqual(ty.layout, layout)
self.assertEqual(ty.dtypes, dtypes)
views = [types.Array(dtype, 0, "C") for dtype in dtypes]
if len(views) > 1:
self.assertEqual(
ty.yield_type,
types.BaseTuple.from_types(views))
else:
self.assertEqual(ty.yield_type, views[0])
if indexers is not None:
self.assertEqual(ty.indexers, indexers)
f32 = types.float32
c64 = types.complex64
i16 = types.int16
a = types.Array(f32, 1, "C")
b = types.Array(f32, 2, "C")
c = types.Array(c64, 2, "F")
d = types.Array(i16, 2, "A")
e = types.Array(i16, 0, "C")
f = types.Array(f32, 1, "A")
g = types.Array(f32, 0, "C")
# 0-dim iterator
ty = types.NumpyNdIterType((e,))
check(ty, (i16,), 0, "C", [('0d', 0, 0, [0])])
self.assertFalse(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((e, g))
check(ty, (i16, f32), 0, "C", [('0d', 0, 0, [0, 1])])
self.assertFalse(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((e, c64))
check(ty, (i16, c64), 0, "C",
[('0d', 0, 0, [0]), ('scalar', 0, 0, [1])])
self.assertFalse(ty.need_shaped_indexing)
# 1-dim iterator
ty = types.NumpyNdIterType((a,))
check(ty, (f32,), 1, "C",
[('flat', 0, 1, [0])])
self.assertFalse(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((a, a))
check(ty, (f32, f32), 1, "C",
[('flat', 0, 1, [0, 1])])
self.assertFalse(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((a, e, e, c64))
check(ty, (f32, i16, i16, c64), 1, "C",
[('flat', 0, 1, [0]), # a
('0d', 0, 0, [1, 2]), # e, e
('scalar', 0, 0, [3]), # c64
])
self.assertFalse(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((a, f))
check(ty, (f32, f32), 1, "C",
[('flat', 0, 1, [0]), ('indexed', 0, 1, [1])])
self.assertTrue(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((f,))
check(ty, (f32,), 1, "C", [('indexed', 0, 1, [0])])
self.assertTrue(ty.need_shaped_indexing)
# 2-dim C-order iterator
ty = types.NumpyNdIterType((b,))
check(ty, (f32,), 2, "C", [('flat', 0, 2, [0])])
self.assertFalse(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((b, c))
check(
ty, (f32, c64), 2, "C", [
('flat', 0, 2, [0]), ('indexed', 0, 2, [1])])
self.assertTrue(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((d,))
check(ty, (i16,), 2, "C", [('indexed', 0, 2, [0])])
self.assertTrue(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((b, c, d, d, e))
check(ty, (f32, c64, i16, i16, i16), 2, "C",
[('flat', 0, 2, [0]), # b
('indexed', 0, 2, [1, 2, 3]), # c, d, d
('0d', 0, 0, [4]), # e
])
self.assertTrue(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((a, b, c, d, d, f))
check(ty, (f32, f32, c64, i16, i16, f32), 2, "C",
[('flat', 1, 2, [0]), # a
('flat', 0, 2, [1]), # b
('indexed', 0, 2, [2, 3, 4]), # c, d, d
('indexed', 1, 2, [5]), # f
])
self.assertTrue(ty.need_shaped_indexing)
# 2-dim F-order iterator
ty = types.NumpyNdIterType((c,))
check(ty, (c64,), 2, "F", [('flat', 0, 2, [0])])
self.assertFalse(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((c, b, c, f))
check(ty, (c64, f32, c64, f32), 2, "F",
[('flat', 0, 2, [0, 2]), # c, c
('indexed', 0, 2, [1]), # b
('indexed', 0, 1, [3]), # f
])
self.assertTrue(ty.need_shaped_indexing)
ty = types.NumpyNdIterType((b, c, c, d, d, a, e))
check(ty, (f32, c64, c64, i16, i16, f32, i16), 2, "F",
[('indexed', 0, 2, [0, 3, 4]), # b, d, d
('flat', 0, 2, [1, 2]), # c, c
('flat', 0, 1, [5]), # a
('0d', 0, 0, [6]), # e
])
self.assertTrue(ty.need_shaped_indexing)
