Repository: numba/numba

Issue:
next(generator) is not supported in nonpython mode
Even though the docs claim that generators are fully supported, one cannot call `next` on them inside another function.

For example:

```
@numba.jit(nopython=True)
def congruential_rng(seed):
    i = seed
    a = 1103515245
    c = 12345
    m = RAND_MAX
    while True:
        #This is a faster shortcut for modulo, 
        #for a power of 2
        i = (a*i + c) & (m)
        yield i

@numba.jit(nopython=True)
def x():
    rng = congruential_rng(0)
    return next(rng)
    #for x in rng:
    #    return x
x()

```

Gives an error:

```
TypingError: Failed at nopython (nopython frontend)
Untyped global name 'next'
File "<ipython-input-51-27e2dbe92485>", line 4
```

Although the commented version with the for loop works.



Discussion:
Thanks for catching this.  We should be able to add this in pretty quickly.

@pitrou: Can you take a look?


Resolving patch:
diff --git a/numba/targets/builtins.py b/numba/targets/builtins.py
--- a/numba/targets/builtins.py
+++ b/numba/targets/builtins.py
@@ -10,7 +10,7 @@
 import llvmlite.llvmpy.core as lc
 
 from .imputils import (lower_builtin, lower_getattr, lower_getattr_generic,
-                       lower_cast, iternext_impl,
+                       lower_cast, iternext_impl, call_getiter, call_iternext,
                        impl_ret_borrowed, impl_ret_untracked)
 from . import optional
 from .. import typing, types, cgutils, utils
@@ -1473,6 +1473,27 @@ def type_impl(context, builder, sig, args):
     return context.get_dummy_value()
 
 
+@lower_builtin(iter, types.IterableType)
+def iter_impl(context, builder, sig, args):
+    ty, = sig.args
+    val, = args
+    iterval = call_getiter(context, builder, ty, val)
+    return iterval
+
+
+@lower_builtin(next, types.IteratorType)
+def next_impl(context, builder, sig, args):
+    iterty, = sig.args
+    iterval, = args
+
+    res = call_iternext(context, builder, iterty, iterval)
+
+    with builder.if_then(builder.not_(res.is_valid()), likely=False):
+        context.call_conv.return_user_exc(builder, StopIteration, ())
+
+    return res.yielded_value()
+
+
 # -----------------------------------------------------------------------------
 
 @lower_builtin("not in", types.Any, types.Any)
diff --git a/numba/typing/builtins.py b/numba/typing/builtins.py
--- a/numba/typing/builtins.py
+++ b/numba/typing/builtins.py
@@ -796,6 +796,29 @@ def generic(self, args, kws):
             zip_type = types.ZipType(args)
             return signature(zip_type, *args)
 
+
+@infer_global(iter)
+class Iter(AbstractTemplate):
+
+    def generic(self, args, kws):
+        assert not kws
+        if len(args) == 1:
+            it = args[0]
+            if isinstance(it, types.IterableType):
+                return signature(it.iterator_type, *args)
+
+
+@infer_global(next)
+class Next(AbstractTemplate):
+
+    def generic(self, args, kws):
+        assert not kws
+        if len(args) == 1:
+            it = args[0]
+            if isinstance(it, types.IteratorType):
+                return signature(it.yield_type, *args)
+
+
 #------------------------------------------------------------------------------
 
 @infer_global(type)


Test patch:
diff --git a/numba/tests/test_array_iterators.py b/numba/tests/test_array_iterators.py
--- a/numba/tests/test_array_iterators.py
+++ b/numba/tests/test_array_iterators.py
@@ -66,6 +66,11 @@ def np_ndindex_array(arr):
             s = s + (i + 1) * (j + 1)
     return s
 
+def iter_next(arr):
+    it = iter(arr)
+    it2 = iter(arr)
+    return next(it), next(it), next(it2)
+
 
 class TestArrayIterators(MemoryLeakMixin, TestCase):
     """
@@ -316,7 +321,7 @@ def test_np_ndindex(self):
     @tag('important')
     def test_np_ndindex_array(self):
         func = np_ndindex_array
-        arr = np.arange(12, dtype=np.int32)
+        arr = np.arange(12, dtype=np.int32) + 10
         self.check_array_unary(arr, typeof(arr), func)
         arr = arr.reshape((4, 3))
         self.check_array_unary(arr, typeof(arr), func)
@@ -329,6 +334,13 @@ def test_np_ndindex_empty(self):
         cfunc = cres.entry_point
         self.assertPreciseEqual(cfunc(), func())
 
+    @tag('important')
+    def test_iter_next(self):
+        # This also checks memory management with iter() and next()
+        func = iter_next
+        arr = np.arange(12, dtype=np.int32) + 10
+        self.check_array_unary(arr, typeof(arr), func)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/numba/tests/test_builtins.py b/numba/tests/test_builtins.py
--- a/numba/tests/test_builtins.py
+++ b/numba/tests/test_builtins.py
@@ -4,11 +4,14 @@
 import functools
 import sys
 
+import numpy as np
+
 import numba.unittest_support as unittest
 from numba.compiler import compile_isolated, Flags
 from numba import jit, typeof, errors, types, utils
 from .support import TestCase, tag
 
+
 enable_pyobj_flags = Flags()
 enable_pyobj_flags.set("enable_pyobject")
 
@@ -85,6 +88,10 @@ def hex_usecase(x):
 def int_usecase(x, base):
     return int(x, base=base)
 
+def iter_next_usecase(x):
+    it = iter(x)
+    return next(it), next(it)
+
 def locals_usecase(x):
     y = 5
     return locals()['y']
@@ -445,6 +452,23 @@ def test_int_npm(self):
         with self.assertTypingError():
             self.test_int(flags=no_pyobj_flags)
 
+    def test_iter_next(self, flags=enable_pyobj_flags):
+        pyfunc = iter_next_usecase
+        cr = compile_isolated(pyfunc, (types.UniTuple(types.int32, 3),),
+                              flags=flags)
+        cfunc = cr.entry_point
+        self.assertPreciseEqual(cfunc((1, 42, 5)), (1, 42))
+
+        cr = compile_isolated(pyfunc, (types.UniTuple(types.int32, 1),),
+                              flags=flags)
+        cfunc = cr.entry_point
+        with self.assertRaises(StopIteration):
+            cfunc((1,))
+
+    @tag('important')
+    def test_iter_next_npm(self):
+        self.test_iter_next(flags=no_pyobj_flags)
+
     def test_locals(self, flags=enable_pyobj_flags):
         pyfunc = locals_usecase
         with self.assertRaises(errors.ForbiddenConstruct):