662 lines
17 KiB
Python
662 lines
17 KiB
Python
|
import numpy as np
|
||
|
|
|
||
|
|
import unittest
|
||
|
|
from numba import jit, njit
|
||
|
|
from numba.core import types
|
||
|
|
from numba.tests.support import TestCase, MemoryLeakMixin
|
||
|
|
from numba.core.datamodel.testing import test_factory
|
||
|
|
|
||
|
|
forceobj_flags = {'nopython': False, 'forceobj': True}
|
||
|
|
nopython_flags = {'nopython': True}
|
||
|
|
|
||
|
|
|
||
|
|
def make_consumer(gen_func):
|
||
|
|
def consumer(x):
|
||
|
|
res = 0.0
|
||
|
|
for y in gen_func(x):
|
||
|
|
res += y
|
||
|
|
return res
|
||
|
|
|
||
|
|
return consumer
|
||
|
|
|
||
|
|
|
||
|
|
def gen1(x):
|
||
|
|
for i in range(x):
|
||
|
|
yield i
|
||
|
|
|
||
|
|
|
||
|
|
def gen2(x):
|
||
|
|
for i in range(x):
|
||
|
|
yield i
|
||
|
|
for j in range(1, 3):
|
||
|
|
yield i + j
|
||
|
|
|
||
|
|
|
||
|
|
def gen3(x):
|
||
|
|
# Polymorphic yield types must be unified
|
||
|
|
yield x
|
||
|
|
yield x + 1.5
|
||
|
|
yield x + 1j
|
||
|
|
|
||
|
|
|
||
|
|
def gen4(x, y, z):
|
||
|
|
for i in range(3):
|
||
|
|
yield z
|
||
|
|
yield y + z
|
||
|
|
return
|
||
|
|
yield x
|
||
|
|
|
||
|
|
|
||
|
|
def gen5():
|
||
|
|
# The bytecode for this generator doesn't contain any YIELD_VALUE
|
||
|
|
# (it's optimized away). We fail typing it, since the yield type
|
||
|
|
# is entirely undefined.
|
||
|
|
if 0:
|
||
|
|
yield 1
|
||
|
|
|
||
|
|
|
||
|
|
def gen6(a, b):
|
||
|
|
# Infinite loop: exercise computation of state variables
|
||
|
|
x = a + 1
|
||
|
|
while True:
|
||
|
|
y = b + 2
|
||
|
|
yield x + y
|
||
|
|
|
||
|
|
|
||
|
|
def gen7(arr):
|
||
|
|
# Array variable in generator state
|
||
|
|
for i in range(arr.size):
|
||
|
|
yield arr[i]
|
||
|
|
|
||
|
|
|
||
|
|
# Optional arguments and boolean state members
|
||
|
|
def gen8(x=1, y=2, b=False):
|
||
|
|
bb = not b
|
||
|
|
yield x
|
||
|
|
if bb:
|
||
|
|
yield y
|
||
|
|
if b:
|
||
|
|
yield x + y
|
||
|
|
|
||
|
|
|
||
|
|
def genobj(x):
|
||
|
|
object()
|
||
|
|
yield x
|
||
|
|
|
||
|
|
|
||
|
|
def return_generator_expr(x):
|
||
|
|
return (i * 2 for i in x)
|
||
|
|
|
||
|
|
|
||
|
|
def gen_ndindex(shape):
|
||
|
|
for ind in np.ndindex(shape):
|
||
|
|
yield ind
|
||
|
|
|
||
|
|
|
||
|
|
def gen_flat(arr):
|
||
|
|
for val in arr.flat:
|
||
|
|
yield val
|
||
|
|
|
||
|
|
|
||
|
|
def gen_ndenumerate(arr):
|
||
|
|
for tup in np.ndenumerate(arr):
|
||
|
|
yield tup
|
||
|
|
|
||
|
|
|
||
|
|
def gen_bool():
|
||
|
|
yield True
|
||
|
|
|
||
|
|
|
||
|
|
def gen_unification_error():
|
||
|
|
yield None
|
||
|
|
yield 1j
|
||
|
|
|
||
|
|
|
||
|
|
def gen_optional_and_type_unification_error():
|
||
|
|
# yields complex and optional(literalint)
|
||
|
|
i = 0
|
||
|
|
yield 1j
|
||
|
|
while True:
|
||
|
|
i = yield i
|
||
|
|
|
||
|
|
|
||
|
|
def gen_changing_tuple_type():
|
||
|
|
# https://github.com/numba/numba/issues/7295
|
||
|
|
yield 1, 2
|
||
|
|
yield 3, 4
|
||
|
|
|
||
|
|
|
||
|
|
def gen_changing_number_type():
|
||
|
|
# additional test for https://github.com/numba/numba/issues/7295
|
||
|
|
yield 1
|
||
|
|
yield 3.5
|
||
|
|
yield 67.8j
|
||
|
|
|
||
|
|
|
||
|
|
class TestGenerators(MemoryLeakMixin, TestCase):
|
||
|
|
def check_generator(self, pygen, cgen):
|
||
|
|
self.assertEqual(next(cgen), next(pygen))
|
||
|
|
# Use list comprehensions to make sure we trash the generator's
|
||
|
|
# former C stack.
|
||
|
|
expected = [x for x in pygen]
|
||
|
|
got = [x for x in cgen]
|
||
|
|
self.assertEqual(expected, got)
|
||
|
|
with self.assertRaises(StopIteration):
|
||
|
|
next(cgen)
|
||
|
|
|
||
|
|
def check_gen1(self, **kwargs):
|
||
|
|
pyfunc = gen1
|
||
|
|
cr = jit((types.int32,), **kwargs)(pyfunc)
|
||
|
|
pygen = pyfunc(8)
|
||
|
|
cgen = cr(8)
|
||
|
|
self.check_generator(pygen, cgen)
|
||
|
|
|
||
|
|
def test_gen1(self):
|
||
|
|
self.check_gen1(**nopython_flags)
|
||
|
|
|
||
|
|
def test_gen1_objmode(self):
|
||
|
|
self.check_gen1(**forceobj_flags)
|
||
|
|
|
||
|
|
def check_gen2(self, **kwargs):
|
||
|
|
pyfunc = gen2
|
||
|
|
cr = jit((types.int32,), **kwargs)(pyfunc)
|
||
|
|
pygen = pyfunc(8)
|
||
|
|
cgen = cr(8)
|
||
|
|
self.check_generator(pygen, cgen)
|
||
|
|
|
||
|
|
def test_gen2(self):
|
||
|
|
self.check_gen2(**nopython_flags)
|
||
|
|
|
||
|
|
def test_gen2_objmode(self):
|
||
|
|
self.check_gen2(**forceobj_flags)
|
||
|
|
|
||
|
|
def check_gen3(self, **kwargs):
|
||
|
|
pyfunc = gen3
|
||
|
|
cr = jit((types.int32,), **kwargs)(pyfunc)
|
||
|
|
pygen = pyfunc(8)
|
||
|
|
cgen = cr(8)
|
||
|
|
self.check_generator(pygen, cgen)
|
||
|
|
|
||
|
|
def test_gen3(self):
|
||
|
|
self.check_gen3(**nopython_flags)
|
||
|
|
|
||
|
|
def test_gen3_objmode(self):
|
||
|
|
self.check_gen3(**forceobj_flags)
|
||
|
|
|
||
|
|
def check_gen4(self, **kwargs):
|
||
|
|
pyfunc = gen4
|
||
|
|
cr = jit((types.int32,) * 3, **kwargs)(pyfunc)
|
||
|
|
pygen = pyfunc(5, 6, 7)
|
||
|
|
cgen = cr(5, 6, 7)
|
||
|
|
self.check_generator(pygen, cgen)
|
||
|
|
|
||
|
|
def test_gen4(self):
|
||
|
|
self.check_gen4(**nopython_flags)
|
||
|
|
|
||
|
|
def test_gen4_objmode(self):
|
||
|
|
self.check_gen4(**forceobj_flags)
|
||
|
|
|
||
|
|
def test_gen5(self):
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
jit((), **nopython_flags)(gen5)
|
||
|
|
self.assertIn("Cannot type generator: it does not yield any value",
|
||
|
|
str(raises.exception))
|
||
|
|
|
||
|
|
def test_gen5_objmode(self):
|
||
|
|
cgen = jit((), **forceobj_flags)(gen5)()
|
||
|
|
self.assertEqual(list(cgen), [])
|
||
|
|
with self.assertRaises(StopIteration):
|
||
|
|
next(cgen)
|
||
|
|
|
||
|
|
def check_gen6(self, **kwargs):
|
||
|
|
cr = jit((types.int32,) * 2, **kwargs)(gen6)
|
||
|
|
cgen = cr(5, 6)
|
||
|
|
l = []
|
||
|
|
for i in range(3):
|
||
|
|
l.append(next(cgen))
|
||
|
|
self.assertEqual(l, [14] * 3)
|
||
|
|
|
||
|
|
def test_gen6(self):
|
||
|
|
self.check_gen6(**nopython_flags)
|
||
|
|
|
||
|
|
def test_gen6_objmode(self):
|
||
|
|
self.check_gen6(**forceobj_flags)
|
||
|
|
|
||
|
|
def check_gen7(self, **kwargs):
|
||
|
|
pyfunc = gen7
|
||
|
|
cr = jit((types.Array(types.float64, 1, 'C'),), **kwargs)(pyfunc)
|
||
|
|
arr = np.linspace(1, 10, 7)
|
||
|
|
pygen = pyfunc(arr.copy())
|
||
|
|
cgen = cr(arr)
|
||
|
|
self.check_generator(pygen, cgen)
|
||
|
|
|
||
|
|
def test_gen7(self):
|
||
|
|
self.check_gen7(**nopython_flags)
|
||
|
|
|
||
|
|
def test_gen7_objmode(self):
|
||
|
|
self.check_gen7(**forceobj_flags)
|
||
|
|
|
||
|
|
def check_gen8(self, **jit_args):
|
||
|
|
pyfunc = gen8
|
||
|
|
cfunc = jit(**jit_args)(pyfunc)
|
||
|
|
|
||
|
|
def check(*args, **kwargs):
|
||
|
|
self.check_generator(pyfunc(*args, **kwargs),
|
||
|
|
cfunc(*args, **kwargs))
|
||
|
|
|
||
|
|
check(2, 3)
|
||
|
|
check(4)
|
||
|
|
check(y=5)
|
||
|
|
check(x=6, b=True)
|
||
|
|
|
||
|
|
def test_gen8(self):
|
||
|
|
self.check_gen8(nopython=True)
|
||
|
|
|
||
|
|
def test_gen8_objmode(self):
|
||
|
|
self.check_gen8(forceobj=True)
|
||
|
|
|
||
|
|
def check_gen9(self, **kwargs):
|
||
|
|
pyfunc = gen_bool
|
||
|
|
cr = jit((), **kwargs)(pyfunc)
|
||
|
|
pygen = pyfunc()
|
||
|
|
cgen = cr()
|
||
|
|
self.check_generator(pygen, cgen)
|
||
|
|
|
||
|
|
def test_gen9(self):
|
||
|
|
self.check_gen9(**nopython_flags)
|
||
|
|
|
||
|
|
def test_gen9_objmode(self):
|
||
|
|
self.check_gen9(**forceobj_flags)
|
||
|
|
|
||
|
|
def check_consume_generator(self, gen_func):
|
||
|
|
cgen = jit(nopython=True)(gen_func)
|
||
|
|
cfunc = jit(nopython=True)(make_consumer(cgen))
|
||
|
|
pyfunc = make_consumer(gen_func)
|
||
|
|
expected = pyfunc(5)
|
||
|
|
got = cfunc(5)
|
||
|
|
self.assertPreciseEqual(got, expected)
|
||
|
|
|
||
|
|
def test_consume_gen1(self):
|
||
|
|
self.check_consume_generator(gen1)
|
||
|
|
|
||
|
|
def test_consume_gen2(self):
|
||
|
|
self.check_consume_generator(gen2)
|
||
|
|
|
||
|
|
def test_consume_gen3(self):
|
||
|
|
self.check_consume_generator(gen3)
|
||
|
|
|
||
|
|
# Check generator storage of some types
|
||
|
|
|
||
|
|
def check_ndindex(self, **kwargs):
|
||
|
|
pyfunc = gen_ndindex
|
||
|
|
cr = jit((types.UniTuple(types.intp, 2),), **kwargs)(pyfunc)
|
||
|
|
shape = (2, 3)
|
||
|
|
pygen = pyfunc(shape)
|
||
|
|
cgen = cr(shape)
|
||
|
|
self.check_generator(pygen, cgen)
|
||
|
|
|
||
|
|
def test_ndindex(self):
|
||
|
|
self.check_ndindex(**nopython_flags)
|
||
|
|
|
||
|
|
def test_ndindex_objmode(self):
|
||
|
|
self.check_ndindex(**forceobj_flags)
|
||
|
|
|
||
|
|
def check_np_flat(self, pyfunc, **kwargs):
|
||
|
|
cr = jit((types.Array(types.int32, 2, "C"),), **kwargs)(pyfunc)
|
||
|
|
arr = np.arange(6, dtype=np.int32).reshape((2, 3))
|
||
|
|
self.check_generator(pyfunc(arr), cr(arr))
|
||
|
|
crA = jit((types.Array(types.int32, 2, "A"),), **kwargs)(pyfunc)
|
||
|
|
arr = arr.T
|
||
|
|
self.check_generator(pyfunc(arr), crA(arr))
|
||
|
|
|
||
|
|
def test_np_flat(self):
|
||
|
|
self.check_np_flat(gen_flat, **nopython_flags)
|
||
|
|
|
||
|
|
def test_np_flat_objmode(self):
|
||
|
|
self.check_np_flat(gen_flat, **forceobj_flags)
|
||
|
|
|
||
|
|
def test_ndenumerate(self):
|
||
|
|
self.check_np_flat(gen_ndenumerate, **nopython_flags)
|
||
|
|
|
||
|
|
def test_ndenumerate_objmode(self):
|
||
|
|
self.check_np_flat(gen_ndenumerate, **forceobj_flags)
|
||
|
|
|
||
|
|
def test_type_unification_error(self):
|
||
|
|
pyfunc = gen_unification_error
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
jit((), **nopython_flags)(pyfunc)
|
||
|
|
|
||
|
|
msg = ("Can't unify yield type from the following types: complex128, "
|
||
|
|
"none")
|
||
|
|
self.assertIn(msg, str(raises.exception))
|
||
|
|
|
||
|
|
def test_optional_expansion_type_unification_error(self):
|
||
|
|
pyfunc = gen_optional_and_type_unification_error
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
jit((), **nopython_flags)(pyfunc)
|
||
|
|
|
||
|
|
msg = ("Can't unify yield type from the following types: complex128, "
|
||
|
|
"int%s, none")
|
||
|
|
self.assertIn(msg % types.intp.bitwidth, str(raises.exception))
|
||
|
|
|
||
|
|
def test_changing_tuple_type(self):
|
||
|
|
# test https://github.com/numba/numba/issues/7295
|
||
|
|
pyfunc = gen_changing_tuple_type
|
||
|
|
expected = list(pyfunc())
|
||
|
|
got = list(njit(pyfunc)())
|
||
|
|
self.assertEqual(expected, got)
|
||
|
|
|
||
|
|
def test_changing_number_type(self):
|
||
|
|
# additional test for https://github.com/numba/numba/issues/7295
|
||
|
|
pyfunc = gen_changing_number_type
|
||
|
|
expected = list(pyfunc())
|
||
|
|
got = list(njit(pyfunc)())
|
||
|
|
self.assertEqual(expected, got)
|
||
|
|
|
||
|
|
|
||
|
|
def nrt_gen0(ary):
|
||
|
|
for elem in ary:
|
||
|
|
yield elem
|
||
|
|
|
||
|
|
|
||
|
|
def nrt_gen1(ary1, ary2):
|
||
|
|
for e1, e2 in zip(ary1, ary2):
|
||
|
|
yield e1
|
||
|
|
yield e2
|
||
|
|
|
||
|
|
|
||
|
|
class TestNrtArrayGen(MemoryLeakMixin, TestCase):
|
||
|
|
def test_nrt_gen0(self):
|
||
|
|
pygen = nrt_gen0
|
||
|
|
cgen = jit(nopython=True)(pygen)
|
||
|
|
|
||
|
|
py_ary = np.arange(10)
|
||
|
|
c_ary = py_ary.copy()
|
||
|
|
|
||
|
|
py_res = list(pygen(py_ary))
|
||
|
|
c_res = list(cgen(c_ary))
|
||
|
|
|
||
|
|
np.testing.assert_equal(py_ary, c_ary)
|
||
|
|
self.assertEqual(py_res, c_res)
|
||
|
|
# Check reference count
|
||
|
|
self.assertRefCountEqual(py_ary, c_ary)
|
||
|
|
|
||
|
|
def test_nrt_gen1(self):
|
||
|
|
pygen = nrt_gen1
|
||
|
|
cgen = jit(nopython=True)(pygen)
|
||
|
|
|
||
|
|
py_ary1 = np.arange(10)
|
||
|
|
py_ary2 = py_ary1 + 100
|
||
|
|
|
||
|
|
c_ary1 = py_ary1.copy()
|
||
|
|
c_ary2 = py_ary2.copy()
|
||
|
|
|
||
|
|
py_res = list(pygen(py_ary1, py_ary2))
|
||
|
|
c_res = list(cgen(c_ary1, c_ary2))
|
||
|
|
|
||
|
|
np.testing.assert_equal(py_ary1, c_ary1)
|
||
|
|
np.testing.assert_equal(py_ary2, c_ary2)
|
||
|
|
self.assertEqual(py_res, c_res)
|
||
|
|
# Check reference count
|
||
|
|
self.assertRefCountEqual(py_ary1, c_ary1)
|
||
|
|
self.assertRefCountEqual(py_ary2, c_ary2)
|
||
|
|
|
||
|
|
def test_combine_gen0_gen1(self):
|
||
|
|
"""
|
||
|
|
Issue #1163 is observed when two generator with NRT object arguments
|
||
|
|
is ran in sequence. The first one does a invalid free and corrupts
|
||
|
|
the NRT memory subsystem. The second generator is likely to segfault
|
||
|
|
due to corrupted NRT data structure (an invalid MemInfo).
|
||
|
|
"""
|
||
|
|
self.test_nrt_gen0()
|
||
|
|
self.test_nrt_gen1()
|
||
|
|
|
||
|
|
def test_nrt_gen0_stop_iteration(self):
|
||
|
|
"""
|
||
|
|
Test cleanup on StopIteration
|
||
|
|
"""
|
||
|
|
pygen = nrt_gen0
|
||
|
|
cgen = jit(nopython=True)(pygen)
|
||
|
|
|
||
|
|
py_ary = np.arange(1)
|
||
|
|
c_ary = py_ary.copy()
|
||
|
|
|
||
|
|
py_iter = pygen(py_ary)
|
||
|
|
c_iter = cgen(c_ary)
|
||
|
|
|
||
|
|
py_res = next(py_iter)
|
||
|
|
c_res = next(c_iter)
|
||
|
|
|
||
|
|
with self.assertRaises(StopIteration):
|
||
|
|
py_res = next(py_iter)
|
||
|
|
|
||
|
|
with self.assertRaises(StopIteration):
|
||
|
|
c_res = next(c_iter)
|
||
|
|
|
||
|
|
del py_iter
|
||
|
|
del c_iter
|
||
|
|
|
||
|
|
np.testing.assert_equal(py_ary, c_ary)
|
||
|
|
self.assertEqual(py_res, c_res)
|
||
|
|
# Check reference count
|
||
|
|
self.assertRefCountEqual(py_ary, c_ary)
|
||
|
|
|
||
|
|
def test_nrt_gen0_no_iter(self):
|
||
|
|
"""
|
||
|
|
Test cleanup for a initialized but never iterated (never call next())
|
||
|
|
generator.
|
||
|
|
"""
|
||
|
|
pygen = nrt_gen0
|
||
|
|
cgen = jit(nopython=True)(pygen)
|
||
|
|
|
||
|
|
py_ary = np.arange(1)
|
||
|
|
c_ary = py_ary.copy()
|
||
|
|
|
||
|
|
py_iter = pygen(py_ary)
|
||
|
|
c_iter = cgen(c_ary)
|
||
|
|
|
||
|
|
del py_iter
|
||
|
|
del c_iter
|
||
|
|
|
||
|
|
np.testing.assert_equal(py_ary, c_ary)
|
||
|
|
|
||
|
|
# Check reference count
|
||
|
|
self.assertRefCountEqual(py_ary, c_ary)
|
||
|
|
|
||
|
|
|
||
|
|
# TODO: fix nested generator and MemoryLeakMixin
|
||
|
|
class TestNrtNestedGen(TestCase):
|
||
|
|
def test_nrt_nested_gen(self):
|
||
|
|
|
||
|
|
def gen0(arr):
|
||
|
|
for i in range(arr.size):
|
||
|
|
yield arr
|
||
|
|
|
||
|
|
def factory(gen0):
|
||
|
|
def gen1(arr):
|
||
|
|
out = np.zeros_like(arr)
|
||
|
|
for x in gen0(arr):
|
||
|
|
out = out + x
|
||
|
|
return out, arr
|
||
|
|
|
||
|
|
return gen1
|
||
|
|
|
||
|
|
py_arr = np.arange(10)
|
||
|
|
c_arr = py_arr.copy()
|
||
|
|
py_res, py_old = factory(gen0)(py_arr)
|
||
|
|
c_gen = jit(nopython=True)(factory(jit(nopython=True)(gen0)))
|
||
|
|
c_res, c_old = c_gen(c_arr)
|
||
|
|
|
||
|
|
self.assertIsNot(py_arr, c_arr)
|
||
|
|
self.assertIs(py_old, py_arr)
|
||
|
|
self.assertIs(c_old, c_arr)
|
||
|
|
|
||
|
|
np.testing.assert_equal(py_res, c_res)
|
||
|
|
|
||
|
|
self.assertRefCountEqual(py_res, c_res)
|
||
|
|
|
||
|
|
# The below test will fail due to generator finalizer not invoked.
|
||
|
|
# This kept a reference of the c_old.
|
||
|
|
#
|
||
|
|
# self.assertEqual(sys.getrefcount(py_old),
|
||
|
|
# sys.getrefcount(c_old))
|
||
|
|
|
||
|
|
@unittest.expectedFailure
|
||
|
|
def test_nrt_nested_gen_refct(self):
|
||
|
|
def gen0(arr):
|
||
|
|
yield arr
|
||
|
|
|
||
|
|
def factory(gen0):
|
||
|
|
def gen1(arr):
|
||
|
|
for out in gen0(arr):
|
||
|
|
return out
|
||
|
|
|
||
|
|
return gen1
|
||
|
|
|
||
|
|
py_arr = np.arange(10)
|
||
|
|
c_arr = py_arr.copy()
|
||
|
|
py_old = factory(gen0)(py_arr)
|
||
|
|
c_gen = jit(nopython=True)(factory(jit(nopython=True)(gen0)))
|
||
|
|
c_old = c_gen(c_arr)
|
||
|
|
|
||
|
|
self.assertIsNot(py_arr, c_arr)
|
||
|
|
self.assertIs(py_old, py_arr)
|
||
|
|
self.assertIs(c_old, c_arr)
|
||
|
|
|
||
|
|
self.assertRefCountEqual(py_old, c_old)
|
||
|
|
|
||
|
|
def test_nrt_nested_nopython_gen(self):
|
||
|
|
"""
|
||
|
|
Test nesting three generators
|
||
|
|
"""
|
||
|
|
|
||
|
|
def factory(decor=lambda x: x):
|
||
|
|
@decor
|
||
|
|
def foo(a, n):
|
||
|
|
for i in range(n):
|
||
|
|
yield a[i]
|
||
|
|
a[i] += i
|
||
|
|
|
||
|
|
@decor
|
||
|
|
def bar(n):
|
||
|
|
a = np.arange(n)
|
||
|
|
for i in foo(a, n):
|
||
|
|
yield i * 2
|
||
|
|
for i in range(a.size):
|
||
|
|
yield a[i]
|
||
|
|
|
||
|
|
@decor
|
||
|
|
def cat(n):
|
||
|
|
for i in bar(n):
|
||
|
|
yield i + i
|
||
|
|
|
||
|
|
return cat
|
||
|
|
|
||
|
|
py_gen = factory()
|
||
|
|
c_gen = factory(jit(nopython=True))
|
||
|
|
|
||
|
|
py_res = list(py_gen(10))
|
||
|
|
c_res = list(c_gen(10))
|
||
|
|
|
||
|
|
self.assertEqual(py_res, c_res)
|
||
|
|
|
||
|
|
|
||
|
|
class TestGeneratorWithNRT(MemoryLeakMixin, TestCase):
|
||
|
|
def test_issue_1254(self):
|
||
|
|
"""
|
||
|
|
Missing environment for returning array
|
||
|
|
"""
|
||
|
|
|
||
|
|
@jit(nopython=True)
|
||
|
|
def random_directions(n):
|
||
|
|
for i in range(n):
|
||
|
|
vec = np.empty(3)
|
||
|
|
vec[:] = 12
|
||
|
|
yield vec
|
||
|
|
|
||
|
|
outputs = list(random_directions(5))
|
||
|
|
self.assertEqual(len(outputs), 5)
|
||
|
|
|
||
|
|
expect = np.empty(3)
|
||
|
|
expect[:] = 12
|
||
|
|
for got in outputs:
|
||
|
|
np.testing.assert_equal(expect, got)
|
||
|
|
|
||
|
|
def test_issue_1265(self):
|
||
|
|
"""
|
||
|
|
Double-free for locally allocated, non escaping NRT objects
|
||
|
|
"""
|
||
|
|
|
||
|
|
def py_gen(rmin, rmax, nr):
|
||
|
|
a = np.linspace(rmin, rmax, nr)
|
||
|
|
yield a[0]
|
||
|
|
yield a[1]
|
||
|
|
|
||
|
|
c_gen = jit(nopython=True)(py_gen)
|
||
|
|
|
||
|
|
py_res = list(py_gen(-2, 2, 100))
|
||
|
|
c_res = list(c_gen(-2, 2, 100))
|
||
|
|
|
||
|
|
self.assertEqual(py_res, c_res)
|
||
|
|
|
||
|
|
def py_driver(args):
|
||
|
|
rmin, rmax, nr = args
|
||
|
|
points = np.empty(nr, dtype=np.complex128)
|
||
|
|
for i, c in enumerate(py_gen(rmin, rmax, nr)):
|
||
|
|
points[i] = c
|
||
|
|
|
||
|
|
return points
|
||
|
|
|
||
|
|
@jit(nopython=True)
|
||
|
|
def c_driver(args):
|
||
|
|
rmin, rmax, nr = args
|
||
|
|
points = np.empty(nr, dtype=np.complex128)
|
||
|
|
for i, c in enumerate(c_gen(rmin, rmax, nr)):
|
||
|
|
points[i] = c
|
||
|
|
|
||
|
|
return points
|
||
|
|
|
||
|
|
n = 2
|
||
|
|
patches = (-2, -1, n)
|
||
|
|
|
||
|
|
py_res = py_driver(patches)
|
||
|
|
# The error will cause a segfault here
|
||
|
|
c_res = c_driver(patches)
|
||
|
|
|
||
|
|
np.testing.assert_equal(py_res, c_res)
|
||
|
|
|
||
|
|
def test_issue_1808(self):
|
||
|
|
"""
|
||
|
|
Incorrect return data model
|
||
|
|
"""
|
||
|
|
magic = 0xdeadbeef
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def generator():
|
||
|
|
yield magic
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def get_generator():
|
||
|
|
return generator()
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def main():
|
||
|
|
out = 0
|
||
|
|
for x in get_generator():
|
||
|
|
out += x
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
self.assertEqual(main(), magic)
|
||
|
|
|
||
|
|
|
||
|
|
class TestGeneratorModel(test_factory()):
|
||
|
|
fe_type = types.Generator(gen_func=None, yield_type=types.int32,
|
||
|
|
arg_types=[types.int64, types.float32],
|
||
|
|
state_types=[types.intp, types.intp[::1]],
|
||
|
|
has_finalizer=False)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|