implement 'async for' loop statement (PEP 492)

Cython/Compiler/Builtin.py

@@ -398,9 +398,16 @@ def init_builtins():
     init_builtin_structs()
     init_builtin_types()
     init_builtin_funcs()
+
     builtin_scope.declare_var(
         '__debug__', PyrexTypes.c_const_type(PyrexTypes.c_bint_type),
         pos=None, cname='(!Py_OptimizeFlag)', is_cdef=True)
+
+    entry = builtin_scope.declare_var(
+        'StopAsyncIteration', PyrexTypes.py_object_type,
+        pos=None, cname='__Pyx_PyExc_StopAsyncIteration')
+    entry.utility_code = UtilityCode.load_cached("StopAsyncIteration", "Coroutine.c")
+
     global list_type, tuple_type, dict_type, set_type, frozenset_type
     global bytes_type, str_type, unicode_type, basestring_type, slice_type
     global float_type, bool_type, type_type, complex_type, bytearray_type

Cython/Compiler/Code.py

@@ -49,7 +49,6 @@ non_portable_builtins_map = {
     'basestring'    : ('PY_MAJOR_VERSION >= 3', 'str'),
     'xrange'        : ('PY_MAJOR_VERSION >= 3', 'range'),
     'raw_input'     : ('PY_MAJOR_VERSION >= 3', 'input'),
-    'StopAsyncIteration': ('PY_VERSION_HEX < 0x030500B1', 'StopIteration'),
 }
 
 basicsize_builtins_map = {

Cython/Compiler/FlowControl.py

@@ -991,6 +991,9 @@ class ControlFlowAnalysis(CythonTransform):
 
             self.mark_assignment(target, node.item)
 
+    def visit_AsyncForStatNode(self, node):
+        return self.visit_ForInStatNode(node)
+
     def visit_ForInStatNode(self, node):
         condition_block = self.flow.nextblock()
         next_block = self.flow.newblock()
@@ -1002,6 +1005,9 @@ class ControlFlowAnalysis(CythonTransform):
 
         if isinstance(node, Nodes.ForInStatNode):
             self.mark_forloop_target(node)
+        elif isinstance(node, Nodes.AsyncForStatNode):
+            # not entirely correct, but good enough for now
+            self.mark_assignment(node.target, node.item)
         else: # Parallel
             self.mark_assignment(node.target)
 

Cython/Compiler/ModuleNode.py

@@ -2071,7 +2071,7 @@ class ModuleNode(Nodes.Node, Nodes.BlockNode):
         code.putln("%s = PyBytes_FromStringAndSize(\"\", 0); %s" % (
             Naming.empty_bytes, code.error_goto_if_null(Naming.empty_bytes, self.pos)))
 
-        for ext_type in ('CyFunction', 'FusedFunction', 'Coroutine', 'Generator'):
+        for ext_type in ('CyFunction', 'FusedFunction', 'Coroutine', 'Generator', 'StopAsyncIteration'):
             code.putln("#ifdef __Pyx_%s_USED" % ext_type)
             code.put_error_if_neg(self.pos, "__pyx_%s_init()" % ext_type)
             code.putln("#endif")

Cython/Compiler/Nodes.py

@@ -6069,40 +6069,49 @@ class DictIterationNextNode(Node):
             target.generate_assignment_code(result, code)
             var.release(code)
 
+
 def ForStatNode(pos, **kw):
     if 'iterator' in kw:
-        return ForInStatNode(pos, **kw)
+        if kw['iterator'].is_async:
+            return AsyncForStatNode(pos, **kw)
+        else:
+            return ForInStatNode(pos, **kw)
     else:
         return ForFromStatNode(pos, **kw)
 
-class ForInStatNode(LoopNode, StatNode):
-    #  for statement
+
+class _ForInStatNode(LoopNode, StatNode):
+    #  Base class of 'for-in' statements.
     #
     #  target        ExprNode
-    #  iterator      IteratorNode
+    #  iterator      IteratorNode | AwaitExprNode(AsyncIteratorNode)
     #  body          StatNode
     #  else_clause   StatNode
-    #  item          NextNode       used internally
+    #  item          NextNode | AwaitExprNode(AsyncNextNode)
+    #  is_async      boolean        true for 'async for' statements
 
-    child_attrs = ["target", "iterator", "body", "else_clause"]
+    child_attrs = ["target", "item", "iterator", "body", "else_clause"]
     item = None
+    is_async = False
+
+    def _create_item_node(self):
+        raise NotImplementedError("must be implemented by subclasses")
 
     def analyse_declarations(self, env):
-        from . import ExprNodes
         self.target.analyse_target_declaration(env)
         self.body.analyse_declarations(env)
         if self.else_clause:
             self.else_clause.analyse_declarations(env)
-        self.item = ExprNodes.NextNode(self.iterator)
+        self._create_item_node()
 
     def analyse_expressions(self, env):
         self.target = self.target.analyse_target_types(env)
         self.iterator = self.iterator.analyse_expressions(env)
-        from . import ExprNodes
-        self.item = ExprNodes.NextNode(self.iterator)  # must rewrap after analysis
+        self._create_item_node()  # must rewrap self.item after analysis
         self.item = self.item.analyse_expressions(env)
-        if (self.iterator.type.is_ptr or self.iterator.type.is_array) and \
-            self.target.type.assignable_from(self.iterator.type):
+        if (not self.is_async and
+                (self.iterator.type.is_ptr or self.iterator.type.is_array) and
+                self.target.type.assignable_from(self.iterator.type)):
             # C array slice optimization.
             pass
         else:
@@ -6168,6 +6177,37 @@ class ForInStatNode(LoopNode, StatNode):
         self.item.annotate(code)
 
 
+class ForInStatNode(_ForInStatNode):
+    #  'for' statement
+
+    is_async = False
+
+    def _create_item_node(self):
+        from .ExprNodes import NextNode
+        self.item = NextNode(self.iterator)
+
+
+class AsyncForStatNode(_ForInStatNode):
+    #  'async for' statement
+    #
+    #  iterator      AwaitExprNode(AsyncIteratorNode)
+    #  item          AwaitIterNextExprNode(AsyncIteratorNode)
+
+    is_async = True
+
+    def __init__(self, pos, iterator, **kw):
+        assert 'item' not in kw
+        from . import ExprNodes
+        # AwaitExprNodes must appear before running MarkClosureVisitor
+        kw['iterator'] = ExprNodes.AwaitExprNode(iterator.pos, arg=iterator)
+        kw['item'] = ExprNodes.AwaitIterNextExprNode(iterator.pos, arg=None)
+        _ForInStatNode.__init__(self, pos, **kw)
+
+    def _create_item_node(self):
+        from . import ExprNodes
+        self.item.arg = ExprNodes.AsyncNextNode(self.iterator)
+
+
 class ForFromStatNode(LoopNode, StatNode):
     #  for name from expr rel name rel expr
     #

Cython/Compiler/Parsing.pxd

@@ -44,7 +44,6 @@ cdef p_typecast(PyrexScanner s)
 cdef p_sizeof(PyrexScanner s)
 cdef p_yield_expression(PyrexScanner s)
 cdef p_yield_statement(PyrexScanner s)
-cdef p_await_expression(PyrexScanner s)
 cdef p_async_statement(PyrexScanner s, ctx)
 cdef p_power(PyrexScanner s)
 cdef p_new_expr(PyrexScanner s)
@@ -109,13 +108,13 @@ cdef p_if_statement(PyrexScanner s)
 cdef p_if_clause(PyrexScanner s)
 cdef p_else_clause(PyrexScanner s)
 cdef p_while_statement(PyrexScanner s)
-cdef p_for_statement(PyrexScanner s)
-cdef dict p_for_bounds(PyrexScanner s, bint allow_testlist = *)
+cdef p_for_statement(PyrexScanner s, bint is_async=*)
+cdef dict p_for_bounds(PyrexScanner s, bint allow_testlist=*, bint is_async=*)
 cdef p_for_from_relation(PyrexScanner s)
 cdef p_for_from_step(PyrexScanner s)
 cdef p_target(PyrexScanner s, terminator)
 cdef p_for_target(PyrexScanner s)
-cdef p_for_iterator(PyrexScanner s, bint allow_testlist = *)
+cdef p_for_iterator(PyrexScanner s, bint allow_testlist=*, bint is_async=*)
 cdef p_try_statement(PyrexScanner s)
 cdef p_except_clause(PyrexScanner s)
 cdef p_include_statement(PyrexScanner s, ctx)

Cython/Compiler/Parsing.py

@@ -390,8 +390,7 @@ def p_async_statement(s, ctx, decorators):
     elif decorators:
         s.error("Decorators can only be followed by functions or classes")
     elif s.sy == 'for':
-        #s.error("'async for' is not currently supported", fatal=False)
-        return p_statement(s, ctx)  # TODO: implement
+        return p_for_statement(s, is_async=True)
     elif s.sy == 'with':
         s.next()
         return p_with_items(s, is_async=True)
@@ -399,10 +398,6 @@ def p_async_statement(s, ctx, decorators):
         s.error("expected one of 'def', 'for', 'with' after 'async'")
 
 
-def p_await_expression(s):
-    n1 = p_atom(s)
-
-
 #power: atom_expr ('**' factor)*
 #atom_expr: ['await'] atom trailer*
 
@@ -1604,23 +1599,25 @@ def p_while_statement(s):
         condition = test, body = body,
         else_clause = else_clause)
 
-def p_for_statement(s):
+
+def p_for_statement(s, is_async=False):
     # s.sy == 'for'
     pos = s.position()
     s.next()
-    kw = p_for_bounds(s, allow_testlist=True)
+    kw = p_for_bounds(s, allow_testlist=True, is_async=is_async)
     body = p_suite(s)
     else_clause = p_else_clause(s)
-    kw.update(body = body, else_clause = else_clause)
+    kw.update(body=body, else_clause=else_clause, is_async=is_async)
     return Nodes.ForStatNode(pos, **kw)
 
-def p_for_bounds(s, allow_testlist=True):
+
+def p_for_bounds(s, allow_testlist=True, is_async=False):
     target = p_for_target(s)
     if s.sy == 'in':
         s.next()
-        iterator = p_for_iterator(s, allow_testlist)
-        return dict( target = target, iterator = iterator )
-    elif not s.in_python_file:
+        iterator = p_for_iterator(s, allow_testlist, is_async=is_async)
+        return dict(target=target, iterator=iterator)
+    elif not s.in_python_file and not is_async:
         if s.sy == 'from':
             s.next()
             bound1 = p_bit_expr(s)
@@ -1690,16 +1687,19 @@ def p_target(s, terminator):
     else:
         return expr
 
+
 def p_for_target(s):
     return p_target(s, 'in')
 
-def p_for_iterator(s, allow_testlist=True):
+
+def p_for_iterator(s, allow_testlist=True, is_async=False):
     pos = s.position()
     if allow_testlist:
         expr = p_testlist(s)
     else:
         expr = p_or_test(s)
-    return ExprNodes.IteratorNode(pos, sequence = expr)
+    return (ExprNodes.AsyncIteratorNode if is_async else ExprNodes.IteratorNode)(pos, sequence=expr)
+
 
 def p_try_statement(s):
     # s.sy == 'try'

Cython/Utility/Coroutine.c

@@ -155,6 +155,56 @@ bad:
 }
 
 
+//////////////////// AsyncIter.proto ////////////////////
+
+static CYTHON_INLINE PyObject *__Pyx_Coroutine_GetAsyncIter(PyObject *o); /*proto*/
+static CYTHON_INLINE PyObject *__Pyx_Coroutine_AsyncIterNext(PyObject *o); /*proto*/
+
+//////////////////// AsyncIter ////////////////////
+//@requires: GetAwaitIter
+//@requires: ObjectHandling.c::PyObjectCallMethod0
+
+static CYTHON_INLINE PyObject *__Pyx_Coroutine_GetAsyncIter(PyObject *obj) {
+#if PY_VERSION_HEX >= 0x030500B1
+    PyAsyncMethods* am = Py_TYPE(obj)->tp_as_async;
+    if (likely(am && am->am_aiter)) {
+        return (*am->am_aiter)(obj);
+    }
+#else
+    PyObject *iter = __Pyx_PyObject_CallMethod0(obj, PYIDENT("__aiter__"));
+    if (likely(iter))
+        return iter;
+    // FIXME: for the sake of a nicely conforming exception message, assume any AttributeError meant '__aiter__'
+    if (!PyErr_ExceptionMatches(PyExc_AttributeError))
+        return NULL;
+#endif
+
+    PyErr_Format(PyExc_TypeError, "'async for' requires an object with __aiter__ method, got %.100s",
+                 Py_TYPE(obj)->tp_name);
+    return NULL;
+}
+
+static CYTHON_INLINE PyObject *__Pyx_Coroutine_AsyncIterNext(PyObject *obj) {
+#if PY_VERSION_HEX >= 0x030500B1
+    PyAsyncMethods* am = Py_TYPE(obj)->tp_as_async;
+    if (likely(am && am->am_anext)) {
+        return (*am->am_anext)(obj);
+    } else {
+#else
+    PyObject *value = __Pyx_PyObject_CallMethod0(obj, PYIDENT("__anext__"));
+    if (likely(value))
+        return value;
+    // FIXME: for the sake of a nicely conforming exception message, assume any AttributeError meant '__anext__'
+    if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
+#endif
+
+        PyErr_Format(PyExc_TypeError, "'async for' requires an object with __anext__ method, got %.100s",
+                     Py_TYPE(obj)->tp_name);
+    }
+    return NULL;
+}
+
+
 //////////////////// pep479.proto ////////////////////
 
 static void __Pyx_Generator_Replace_StopIteration(void); /*proto*/
@@ -1531,3 +1581,102 @@ except AttributeError:
 #endif
     return module;
 }
+
+
+//////////////////// StopAsyncIteration.proto ////////////////////
+
+#define __Pyx_StopAsyncIteration_USED
+static PyObject *__Pyx_PyExc_StopAsyncIteration;
+static int __pyx_StopAsyncIteration_init(void); /*proto*/
+
+//////////////////// StopAsyncIteration ////////////////////
+
+#if PY_VERSION_HEX < 0x030500B1
+static PyTypeObject __Pyx__PyExc_StopAsyncIteration_type = {
+    PyVarObject_HEAD_INIT(0, 0)
+    "StopAsyncIteration",               /*tp_name*/
+    sizeof(PyBaseExceptionObject),      /*tp_basicsize*/
+    0,                                  /*tp_itemsize*/
+    0,                                  /*tp_dealloc*/
+    0,                                  /*tp_print*/
+    0,                                  /*tp_getattr*/
+    0,                                  /*tp_setattr*/
+#if PY_MAJOR_VERSION < 3
+    0,                                  /*tp_compare*/
+#else
+    0,                                  /*reserved*/
+#endif
+    0,                                  /*tp_repr*/
+    0,                                  /*tp_as_number*/
+    0,                                  /*tp_as_sequence*/
+    0,                                  /*tp_as_mapping*/
+    0,                                  /*tp_hash*/
+    0,                                  /*tp_call*/
+    0,                                  /*tp_str*/
+    0,                                  /*tp_getattro*/
+    0,                                  /*tp_setattro*/
+    0,                                  /*tp_as_buffer*/
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,  /*tp_flags*/
+    PyDoc_STR("Signal the end from iterator.__anext__()."),  /*tp_doc*/
+    0,                                  /*tp_traverse*/
+    0,                                  /*tp_clear*/
+    0,                                  /*tp_richcompare*/
+    0,                                  /*tp_weaklistoffset*/
+    0,                                  /*tp_iter*/
+    0,                                  /*tp_iternext*/
+    0,                                  /*tp_methods*/
+    0,                                  /*tp_members*/
+    0,                                  /*tp_getset*/
+    0,                                  /*tp_base*/
+    0,                                  /*tp_dict*/
+    0,                                  /*tp_descr_get*/
+    0,                                  /*tp_descr_set*/
+    0,                                  /*tp_dictoffset*/
+    0,                                  /*tp_init*/
+    0,                                  /*tp_alloc*/
+    0,                                  /*tp_new*/
+    0,                                  /*tp_free*/
+    0,                                  /*tp_is_gc*/
+    0,                                  /*tp_bases*/
+    0,                                  /*tp_mro*/
+    0,                                  /*tp_cache*/
+    0,                                  /*tp_subclasses*/
+    0,                                  /*tp_weaklist*/
+    0,                                  /*tp_del*/
+    0,                                  /*tp_version_tag*/
+#if PY_VERSION_HEX >= 0x030400a1
+    0,                                  /*tp_finalize*/
+#endif
+};
+#endif
+
+static int __pyx_StopAsyncIteration_init(void) {
+#if PY_VERSION_HEX >= 0x030500B1
+    __Pyx_PyExc_StopAsyncIteration = PyExc_StopAsyncIteration;
+#else
+    PyObject *builtins = PyEval_GetBuiltins();
+    if (likely(builtins)) {
+        PyObject *exc = PyMapping_GetItemString(builtins, "StopAsyncIteration");
+        if (exc) {
+            __Pyx_PyExc_StopAsyncIteration = exc;
+            return 0;
+        }
+    }
+    PyErr_Clear();
+
+    __Pyx__PyExc_StopAsyncIteration_type.tp_traverse = ((PyTypeObject*)PyExc_BaseException)->tp_traverse;
+    __Pyx__PyExc_StopAsyncIteration_type.tp_clear = ((PyTypeObject*)PyExc_BaseException)->tp_clear;
+    __Pyx__PyExc_StopAsyncIteration_type.tp_dealloc = ((PyTypeObject*)PyExc_BaseException)->tp_dealloc;
+    __Pyx__PyExc_StopAsyncIteration_type.tp_dictoffset = ((PyTypeObject*)PyExc_BaseException)->tp_dictoffset;
+    __Pyx__PyExc_StopAsyncIteration_type.tp_base = (PyTypeObject*)PyExc_Exception;
+    __Pyx__PyExc_StopAsyncIteration_type.tp_init = ((PyTypeObject*)PyExc_BaseException)->tp_init;
+    __Pyx__PyExc_StopAsyncIteration_type.tp_new = ((PyTypeObject*)PyExc_BaseException)->tp_new;
+
+    __Pyx_PyExc_StopAsyncIteration = (PyObject*) __Pyx_FetchCommonType(&__Pyx__PyExc_StopAsyncIteration_type);
+    if (unlikely(!__Pyx_PyExc_StopAsyncIteration))
+        return -1;
+    if (builtins && unlikely(PyMapping_SetItemString(builtins, "StopAsyncIteration", __Pyx_PyExc_StopAsyncIteration) < 0))
+        return -1;
+#endif
+    return 0;
+}

tests/run/test_coroutines_pep492.pyx

@@ -104,6 +104,7 @@ class TokenizerRegrTest(unittest.TestCase):
 
 class CoroutineTest(unittest.TestCase):
 
+    @classmethod
     def setUpClass(cls):
         # never mark warnings as "already seen" to prevent them from being suppressed
         from warnings import simplefilter
@@ -209,8 +210,9 @@ class CoroutineTest(unittest.TestCase):
         with check():
             iter(foo())
 
-        with check():
-            next(foo())
+        # in Cython: not iterable, but an iterator ...
+        #with check():
+        #    next(foo())
 
         with silence_coro_gc(), check():
             for i in foo():