Unicode type iteration

numba/dictobject.py

@@ -1126,7 +1126,7 @@ def impl_dict_getiter(context, builder, sig, args):
 
 
 @lower_builtin('iternext', types.DictIteratorType)
-@iternext_impl
+@iternext_impl()
 def impl_iterator_iternext(context, builder, sig, args, result):
     iter_type = sig.args[0]
     it = context.make_helper(builder, iter_type, args[0])

numba/runtime/nrt.c

@@ -295,14 +295,14 @@ void NRT_MemInfo_destroy(NRT_MemInfo *mi) {
 }
 
 void NRT_MemInfo_acquire(NRT_MemInfo *mi) {
-    NRT_Debug(nrt_debug_print("NRT_acquire %p refct=%zu\n", mi,
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_acquire %p refct=%zu\n", mi,
                                                             mi->refct));
     assert(mi->refct > 0 && "RefCt cannot be zero");
     TheMSys.atomic_inc(&mi->refct);
 }
 
 void NRT_MemInfo_call_dtor(NRT_MemInfo *mi) {
-    NRT_Debug(nrt_debug_print("nrt_meminfo_call_dtor %p\n", mi));
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_call_dtor %p\n", mi));
     if (mi->dtor && !TheMSys.shutting)
         /* We have a destructor and the system is not shutting down */
         mi->dtor(mi->data, mi->size, mi->dtor_info);
@@ -311,7 +311,7 @@ void NRT_MemInfo_call_dtor(NRT_MemInfo *mi) {
 }
 
 void NRT_MemInfo_release(NRT_MemInfo *mi) {
-    NRT_Debug(nrt_debug_print("NRT_release %p refct=%zu\n", mi,
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_release %p refct=%zu\n", mi,
                                                             mi->refct));
     assert (mi->refct > 0 && "RefCt cannot be 0");
     /* RefCt drop to zero */
@@ -339,7 +339,7 @@ void NRT_MemInfo_dump(NRT_MemInfo *mi, FILE *out) {
 
 static void
 nrt_varsize_dtor(void *ptr, size_t size, void *info) {
-    NRT_Debug(nrt_debug_print("nrt_buffer_dtor %p\n", ptr));
+    NRT_Debug(nrt_debug_print("nrt_varsize_dtor %p\n", ptr));
     if (info) {
         /* call element dtor */
         typedef void dtor_fn_t(void *ptr);

numba/targets/arrayobj.py

@@ -277,7 +277,7 @@ def _getitem_array1d(context, builder, arrayty, array, idx, wraparound):
     return load_item(context, builder, arrayty, ptr)
 
 @lower_builtin('iternext', types.ArrayIterator)
-@iternext_impl
+@iternext_impl()
 def iternext_array(context, builder, sig, args, result):
     [iterty] = sig.args
     [iter] = args
@@ -2976,7 +2976,7 @@ def make_array_flatiter(context, builder, arrty, arr):
 
 
 @lower_builtin('iternext', types.NumpyFlatType)
-@iternext_impl
+@iternext_impl()
 def iternext_numpy_flatiter(context, builder, sig, args, result):
     [flatiterty] = sig.args
     [flatiter] = args
@@ -3054,7 +3054,7 @@ def make_array_ndenumerate(context, builder, sig, args):
 
 
 @lower_builtin('iternext', types.NumpyNdEnumerateType)
-@iternext_impl
+@iternext_impl()
 def iternext_numpy_nditer(context, builder, sig, args, result):
     [nditerty] = sig.args
     [nditer] = args
@@ -3106,7 +3106,7 @@ def make_array_ndindex(context, builder, sig, args):
     return impl_ret_borrowed(context, builder, sig.return_type, res)
 
 @lower_builtin('iternext', types.NumpyNdIndexType)
-@iternext_impl
+@iternext_impl()
 def iternext_numpy_ndindex(context, builder, sig, args, result):
     [nditerty] = sig.args
     [nditer] = args
@@ -3137,7 +3137,7 @@ def make_array_nditer(context, builder, sig, args):
     return impl_ret_borrowed(context, builder, nditerty, res)
 
 @lower_builtin('iternext', types.NumpyNdIterType)
-@iternext_impl
+@iternext_impl()
 def iternext_numpy_ndindex(context, builder, sig, args, result):
     [nditerty] = sig.args
     [nditer] = args

numba/targets/imputils.py

@@ -234,7 +234,7 @@ def wrapper(cls):
         # These are unbound methods
         iternext = cls.iternext
 
-        @iternext_impl
+        @iternext_impl()
         def iternext_wrapper(context, builder, sig, args, result):
             (value,) = args
             iterobj = cls(context, builder, value)
@@ -294,24 +294,31 @@ def yielded_value(self):
         return self._pairobj.first
 
 
-def iternext_impl(func):
+def iternext_impl(new_ref=False):
     """
     Wrap the given iternext() implementation so that it gets passed
     an _IternextResult() object easing the returning of the iternext()
     result pair.
 
+    new_ref: indicates whether the yielded item is a new ref or a borrowed
+
     The wrapped function will be called with the following signature:
         (context, builder, sig, args, iternext_result)
     """
-
-    def wrapper(context, builder, sig, args):
-        pair_type = sig.return_type
-        pairobj = context.make_helper(builder, pair_type)
-        func(context, builder, sig, args,
-             _IternextResult(context, builder, pairobj))
-        return impl_ret_borrowed(context, builder,
-                                 pair_type, pairobj._getvalue())
-    return wrapper
+    def outer(func):
+        def wrapper(context, builder, sig, args):
+            pair_type = sig.return_type
+            pairobj = context.make_helper(builder, pair_type)
+            func(context, builder, sig, args,
+                _IternextResult(context, builder, pairobj))
+            if new_ref:
+                impl_ret = impl_ret_new_ref
+            else:
+                impl_ret = impl_ret_borrowed
+            return impl_ret(context, builder,
+                                    pair_type, pairobj._getvalue())
+        return wrapper
+    return outer
 
 
 def call_getiter(context, builder, iterable_type, val):

numba/targets/iterators.py

@@ -44,7 +44,7 @@ def make_enumerate_object(context, builder, sig, args):
     return impl_ret_new_ref(context, builder, sig.return_type, res)
 
 @lower_builtin('iternext', types.EnumerateType)
-@iternext_impl
+@iternext_impl()
 def iternext_enumerate(context, builder, sig, args, result):
     [enumty] = sig.args
     [enum] = args
@@ -87,7 +87,7 @@ def make_zip_object(context, builder, sig, args):
     return impl_ret_new_ref(context, builder, sig.return_type, res)
 
 @lower_builtin('iternext', types.ZipType)
-@iternext_impl
+@iternext_impl()
 def iternext_zip(context, builder, sig, args, result):
     [zip_type] = sig.args
     [zipobj] = args
@@ -125,7 +125,7 @@ def iternext_zip(context, builder, sig, args, result):
 # generator implementation
 
 @lower_builtin('iternext', types.Generator)
-@iternext_impl
+@iternext_impl()
 def iternext_zip(context, builder, sig, args, result):
     genty, = sig.args
     gen, = args

numba/targets/listobj.py

@@ -487,7 +487,7 @@ def getiter_list(context, builder, sig, args):
     return impl_ret_borrowed(context, builder, sig.return_type, inst.value)
 
 @lower_builtin('iternext', types.ListIter)
-@iternext_impl
+@iternext_impl()
 def iternext_listiter(context, builder, sig, args, result):
     inst = ListIterInstance(context, builder, sig.args[0], args[0])
 

numba/targets/setobj.py

@@ -1215,7 +1215,7 @@ def getiter_set(context, builder, sig, args):
     return impl_ret_borrowed(context, builder, sig.return_type, inst.value)
 
 @lower_builtin('iternext', types.SetIter)
-@iternext_impl
+@iternext_impl()
 def iternext_listiter(context, builder, sig, args, result):
     inst = SetIterInstance(context, builder, sig.args[0], args[0])
     inst.iternext(result)

numba/targets/tupleobj.py

@@ -148,7 +148,7 @@ def getiter_unituple(context, builder, sig, args):
 
 
 @lower_builtin('iternext', types.UniTupleIter)
-@iternext_impl
+@iternext_impl()
 def iternext_unituple(context, builder, sig, args, result):
     [tupiterty] = sig.args
     [tupiter] = args

numba/tests/test_unicode.py

@@ -104,6 +104,44 @@ def join_empty_usecase(x):
     return x.join(l)
 
 
+def iter_usecase(x):
+    l = []
+    for i in x:
+        l.append(i)
+    return l
+
+
+def literal_iter_usecase():
+    l = []
+    for i in '大处着眼，小处着手。':
+        l.append(i)
+    return l
+
+
+def enumerated_iter_usecase(x):
+    buf = ""
+    scan = 0
+    for i, s in enumerate(x):
+        buf += s
+        scan += 1
+    return buf, scan
+
+
+def iter_stopiteration_usecase(x):
+    n = len(x)
+    i = iter(x)
+    for _ in range(n + 1):
+        next(i)
+
+
+def literal_iter_stopiteration_usecase():
+    s = '大处着眼，小处着手。'
+    i = iter(s)
+    n = len(s)
+    for _ in range(n + 1):
+        next(i)
+
+
 class BaseTest(MemoryLeakMixin, TestCase):
     def setUp(self):
         super(BaseTest, self).setUp()
@@ -333,7 +371,7 @@ def test_split_exception_empty_sep(self):
         # Handle empty separator exception
         for func in [pyfunc, cfunc]:
             with self.assertRaises(ValueError) as raises:
-                    func('a', '')
+                func('a', '')
             self.assertIn('empty separator', str(raises.exception))
 
     def test_split_exception_noninteger_maxsplit(self):
@@ -343,7 +381,7 @@ def test_split_exception_noninteger_maxsplit(self):
         # Handle non-integer maxsplit exception
         for sep in [' ', None]:
             with self.assertRaises(TypingError) as raises:
-                    cfunc('a', sep, 2.4)
+                cfunc('a', sep, 2.4)
             self.assertIn('float64', str(raises.exception),
                           'non-integer maxsplit with sep = %s' % sep)
 
@@ -397,7 +435,7 @@ def test_split_whitespace(self):
         pyfunc = split_whitespace_usecase
         cfunc = njit(pyfunc)
 
-        #list copied from https://github.com/python/cpython/blob/master/Objects/unicodetype_db.h
+        # list copied from https://github.com/python/cpython/blob/master/Objects/unicodetype_db.h
         all_whitespace = ''.join(map(chr, [
             0x0009, 0x000A, 0x000B, 0x000C, 0x000D, 0x001C, 0x001D, 0x001E, 0x001F, 0x0020,
             0x0085, 0x00A0, 0x1680, 0x2000, 0x2001, 0x2002, 0x2003, 0x2004, 0x2005, 0x2006,
@@ -441,12 +479,13 @@ def test_join_non_string_exception(self):
 
         # Handle empty separator exception
         with self.assertRaises(TypingError) as raises:
-            cfunc('', [1,2,3])
+            cfunc('', [1, 2, 3])
         # This error message is obscure, but indicates the error was trapped in typing of str.join()
         # Feel free to change this as we update error messages.
         exc_message = str(raises.exception)
         self.assertIn("Invalid use of BoundFunction", exc_message)
-        self.assertIn("(reflected list(int", exc_message)  # could be int32 or int64
+        # could be int32 or int64
+        self.assertIn("(reflected list(int", exc_message)
 
     def test_join(self):
         pyfunc = join_usecase
@@ -664,5 +703,43 @@ def f():
         self.assertEqual(f.py_func(), f())
 
 
+@unittest.skipUnless(_py34_or_later,
+                     'unicode support requires Python 3.4 or later')
+class TestUnicodeIteration(BaseTest):
+
+    def test_unicode_iter(self):
+        pyfunc = iter_usecase
+        cfunc = njit(pyfunc)
+        for a in UNICODE_EXAMPLES:
+            self.assertPreciseEqual(pyfunc(a), cfunc(a))
+
+    def test_unicode_literal_iter(self):
+        pyfunc = literal_iter_usecase
+        cfunc = njit(pyfunc)
+        self.assertPreciseEqual(pyfunc(), cfunc())
+
+    def test_unicode_enumerate_iter(self):
+        pyfunc = enumerated_iter_usecase
+        cfunc = njit(pyfunc)
+        for a in UNICODE_EXAMPLES:
+            self.assertPreciseEqual(pyfunc(a), cfunc(a))
+
+    def test_unicode_stopiteration_iter(self):
+        self.disable_leak_check()
+        pyfunc = iter_stopiteration_usecase
+        cfunc = njit(pyfunc)
+        for f in (pyfunc, cfunc):
+            for a in UNICODE_EXAMPLES:
+                with self.assertRaises(StopIteration):
+                    f(a)
+
+    def test_unicode_literal_stopiteration_iter(self):
+        pyfunc = literal_iter_stopiteration_usecase
+        cfunc = njit(pyfunc)
+        for f in (pyfunc, cfunc):
+            with self.assertRaises(StopIteration):
+                f()
+
+
 if __name__ == '__main__':
     unittest.main()

numba/types/misc.py

@@ -474,7 +474,25 @@ def get_call_type(self, context, args, kws):
         return typing.signature(self, *posargs)
 
 
-class UnicodeType(Type):
+class UnicodeType(IterableType):
+
     def __init__(self, name):
+        # this is updated on first call to `.iterator_type`.
+        self._iterator_type = None
         super(UnicodeType, self).__init__(name)
 
+    @property
+    def iterator_type(self):
+        # lazily grab the iterator_type, self needs to init before the
+        # iterator_type is instantiated else variety of pickle failures
+        if self._iterator_type is None:
+            self._iterator_type = UnicodeIteratorType(self)
+        return self._iterator_type
+
+
+class UnicodeIteratorType(SimpleIteratorType):
+
+    def __init__(self, dtype):
+        name = "iter_unicode"
+        self.data = dtype
+        super(UnicodeIteratorType, self).__init__(name, dtype)