Issue #28008: Implement PEP 530 -- asynchronous comprehensions.

Grammar/Grammar

@@ -146,7 +146,7 @@ argument: ( test [comp_for] |
             '*' test )
 
 comp_iter: comp_for | comp_if
-comp_for: 'for' exprlist 'in' or_test [comp_iter]
+comp_for: [ASYNC] 'for' exprlist 'in' or_test [comp_iter]
 comp_if: 'if' test_nocond [comp_iter]
 
 # not used in grammar, but may appear in "node" passed from Parser to Compiler

Include/Python-ast.h

@@ -389,6 +389,7 @@ struct _comprehension {
     expr_ty target;
     expr_ty iter;
     asdl_seq *ifs;
+    int is_async;
 };
 
 enum _excepthandler_kind {ExceptHandler_kind=1};
@@ -609,9 +610,9 @@ slice_ty _Py_Slice(expr_ty lower, expr_ty upper, expr_ty step, PyArena *arena);
 slice_ty _Py_ExtSlice(asdl_seq * dims, PyArena *arena);
 #define Index(a0, a1) _Py_Index(a0, a1)
 slice_ty _Py_Index(expr_ty value, PyArena *arena);
-#define comprehension(a0, a1, a2, a3) _Py_comprehension(a0, a1, a2, a3)
+#define comprehension(a0, a1, a2, a3, a4) _Py_comprehension(a0, a1, a2, a3, a4)
 comprehension_ty _Py_comprehension(expr_ty target, expr_ty iter, asdl_seq *
-                                   ifs, PyArena *arena);
+                                   ifs, int is_async, PyArena *arena);
 #define ExceptHandler(a0, a1, a2, a3, a4, a5) _Py_ExceptHandler(a0, a1, a2, a3, a4, a5)
 excepthandler_ty _Py_ExceptHandler(expr_ty type, identifier name, asdl_seq *
                                    body, int lineno, int col_offset, PyArena

Lib/lib2to3/Grammar.txt

@@ -150,7 +150,7 @@ arglist: (argument ',')* (argument [',']
 argument: test [comp_for] | test '=' test  # Really [keyword '='] test
 
 comp_iter: comp_for | comp_if
-comp_for: 'for' exprlist 'in' testlist_safe [comp_iter]
+comp_for: [ASYNC] 'for' exprlist 'in' testlist_safe [comp_iter]
 comp_if: 'if' old_test [comp_iter]
 
 testlist1: test (',' test)*

Lib/lib2to3/tests/test_parser.py

@@ -133,6 +133,24 @@ class TestAsyncAwait(GrammarTest):
                              await x
                       """)
 
+        self.validate("""async def foo():
+                             [i async for i in b]
+                      """)
+
+        self.validate("""async def foo():
+                             {i for i in b
+                                async for i in a if await i
+                                  for b in i}
+                      """)
+
+        self.validate("""async def foo():
+                             [await i for i in b if await c]
+                      """)
+
+        self.validate("""async def foo():
+                             [ i for i in b if c]
+                      """)
+
         self.validate("""async def foo():
 
             def foo(): pass

Lib/test/badsyntax_async1.py

-async def foo(a=await something()):
-    pass

Lib/test/badsyntax_async2.py

-async def foo(a=await something()):
-    pass

Lib/test/badsyntax_async3.py

-async def foo():
-    [i async for i in els]

Lib/test/badsyntax_async4.py

-async def foo():
-    await

Lib/test/badsyntax_async5.py

-def foo():
-    await something()

Lib/test/badsyntax_async7.py

-async def foo():
-    yield from []

Lib/test/badsyntax_async8.py

-async def foo():
-    await await fut

Lib/test/test_coroutines.py

@@ -69,49 +69,130 @@ def silence_coro_gc():
 class AsyncBadSyntaxTest(unittest.TestCase):
 
     def test_badsyntax_1(self):
-        with self.assertRaisesRegex(SyntaxError, "'await' outside"):
-            import test.badsyntax_async1
+        samples = [
+            """def foo():
+                await something()
+            """,
 
-    def test_badsyntax_2(self):
-        with self.assertRaisesRegex(SyntaxError, "'await' outside"):
-            import test.badsyntax_async2
+            """await something()""",
 
-    def test_badsyntax_3(self):
-        with self.assertRaisesRegex(SyntaxError, 'invalid syntax'):
-            import test.badsyntax_async3
+            """async def foo():
+                yield from []
+            """,
 
-    def test_badsyntax_4(self):
-        with self.assertRaisesRegex(SyntaxError, 'invalid syntax'):
-            import test.badsyntax_async4
+            """async def foo():
+                await await fut
+            """,
 
-    def test_badsyntax_5(self):
-        with self.assertRaisesRegex(SyntaxError, 'invalid syntax'):
-            import test.badsyntax_async5
+            """async def foo(a=await something()):
+                pass
+            """,
 
-    def test_badsyntax_7(self):
-        with self.assertRaisesRegex(
-            SyntaxError, "'yield from' inside async function"):
+            """async def foo(a:await something()):
+                pass
+            """,
+
+            """async def foo():
+                def bar():
+                 [i async for i in els]
+            """,
 
-            import test.badsyntax_async7
+            """async def foo():
+                def bar():
+                 [await i for i in els]
+            """,
 
-    def test_badsyntax_8(self):
-        with self.assertRaisesRegex(SyntaxError, 'invalid syntax'):
-            import test.badsyntax_async8
+            """async def foo():
+                def bar():
+                 [i for i in els
+                    async for b in els]
+            """,
 
-    def test_badsyntax_9(self):
-        ns = {}
-        for comp in {'(await a for a in b)',
-                     '[await a for a in b]',
-                     '{await a for a in b}',
-                     '{await a: c for a in b}'}:
+            """async def foo():
+                def bar():
+                 [i for i in els
+                    for c in b
+                    async for b in els]
+            """,
 
-            with self.assertRaisesRegex(SyntaxError, 'await.*in comprehen'):
-                exec('async def f():\n\t{}'.format(comp), ns, ns)
+            """async def foo():
+                def bar():
+                 [i for i in els
+                    async for b in els
+                    for c in b]
+            """,
 
-    def test_badsyntax_10(self):
-        # Tests for issue 24619
+            """async def foo():
+                def bar():
+                 [i for i in els
+                    for b in await els]
+            """,
+
+            """async def foo():
+                def bar():
+                 [i for i in els
+                    for b in els
+                        if await b]
+            """,
+
+            """async def foo():
+                def bar():
+                 [i for i in await els]
+            """,
+
+            """async def foo():
+                def bar():
+                 [i for i in els if await i]
+            """,
+
+            """def bar():
+                 [i async for i in els]
+            """,
+
+            """def bar():
+                 [await i for i in els]
+            """,
+
+            """def bar():
+                 [i for i in els
+                    async for b in els]
+            """,
+
+            """def bar():
+                 [i for i in els
+                    for c in b
+                    async for b in els]
+            """,
+
+            """def bar():
+                 [i for i in els
+                    async for b in els
+                    for c in b]
+            """,
+
+            """def bar():
+                 [i for i in els
+                    for b in await els]
+            """,
+
+            """def bar():
+                 [i for i in els
+                    for b in els
+                        if await b]
+            """,
+
+            """def bar():
+                 [i for i in await els]
+            """,
+
+            """def bar():
+                 [i for i in els if await i]
+            """,
+
+            """async def foo():
+                await
+            """,
 
-        samples = [
             """async def foo():
                    def bar(): pass
                    await = 1
@@ -1531,6 +1612,185 @@ class CoroutineTest(unittest.TestCase):
                 warnings.simplefilter("error")
                 run_async(foo())
 
+    def test_comp_1(self):
+        async def f(i):
+            return i
+
+        async def run_list():
+            return [await c for c in [f(1), f(41)]]
+
+        async def run_set():
+            return {await c for c in [f(1), f(41)]}
+
+        async def run_dict1():
+            return {await c: 'a' for c in [f(1), f(41)]}
+
+        async def run_dict2():
+            return {i: await c for i, c in enumerate([f(1), f(41)])}
+
+        self.assertEqual(run_async(run_list()), ([], [1, 41]))
+        self.assertEqual(run_async(run_set()), ([], {1, 41}))
+        self.assertEqual(run_async(run_dict1()), ([], {1: 'a', 41: 'a'}))
+        self.assertEqual(run_async(run_dict2()), ([], {0: 1, 1: 41}))
+
+    def test_comp_2(self):
+        async def f(i):
+            return i
+
+        async def run_list():
+            return [s for c in [f(''), f('abc'), f(''), f(['de', 'fg'])]
+                    for s in await c]
+
+        self.assertEqual(
+            run_async(run_list()),
+            ([], ['a', 'b', 'c', 'de', 'fg']))
+
+        async def run_set():
+            return {d
+                    for c in [f([f([10, 30]),
+                                 f([20])])]
+                    for s in await c
+                    for d in await s}
+
+        self.assertEqual(
+            run_async(run_set()),
+            ([], {10, 20, 30}))
+
+        async def run_set2():
+            return {await s
+                    for c in [f([f(10), f(20)])]
+                    for s in await c}
+
+        self.assertEqual(
+            run_async(run_set2()),
+            ([], {10, 20}))
+
+    def test_comp_3(self):
+        async def f(it):
+            for i in it:
+                yield i
+
+        async def run_list():
+            return [i + 1 async for i in f([10, 20])]
+        self.assertEqual(
+            run_async(run_list()),
+            ([], [11, 21]))
+
+        async def run_set():
+            return {i + 1 async for i in f([10, 20])}
+        self.assertEqual(
+            run_async(run_set()),
+            ([], {11, 21}))
+
+        async def run_dict():
+            return {i + 1: i + 2 async for i in f([10, 20])}
+        self.assertEqual(
+            run_async(run_dict()),
+            ([], {11: 12, 21: 22}))
+
+        async def run_gen():
+            gen = (i + 1 async for i in f([10, 20]))
+            return [g + 100 async for g in gen]
+        self.assertEqual(
+            run_async(run_gen()),
+            ([], [111, 121]))
+
+    def test_comp_4(self):
+        async def f(it):
+            for i in it:
+                yield i
+
+        async def run_list():
+            return [i + 1 async for i in f([10, 20]) if i > 10]
+        self.assertEqual(
+            run_async(run_list()),
+            ([], [21]))
+
+        async def run_set():
+            return {i + 1 async for i in f([10, 20]) if i > 10}
+        self.assertEqual(
+            run_async(run_set()),
+            ([], {21}))
+
+        async def run_dict():
+            return {i + 1: i + 2 async for i in f([10, 20]) if i > 10}
+        self.assertEqual(
+            run_async(run_dict()),
+            ([], {21: 22}))
+
+        async def run_gen():
+            gen = (i + 1 async for i in f([10, 20]) if i > 10)
+            return [g + 100 async for g in gen]
+        self.assertEqual(
+            run_async(run_gen()),
+            ([], [121]))
+
+    def test_comp_5(self):
+        async def f(it):
+            for i in it:
+                yield i
+
+        async def run_list():
+            return [i + 1 for pair in ([10, 20], [30, 40]) if pair[0] > 10
+                    async for i in f(pair) if i > 30]
+        self.assertEqual(
+            run_async(run_list()),
+            ([], [41]))
+
+    def test_comp_6(self):
+        async def f(it):
+            for i in it:
+                yield i
+
+        async def run_list():
+            return [i + 1 async for seq in f([(10, 20), (30,)])
+                    for i in seq]
+
+        self.assertEqual(
+            run_async(run_list()),
+            ([], [11, 21, 31]))
+
+    def test_comp_7(self):
+        async def f():
+            yield 1
+            yield 2
+            raise Exception('aaa')
+
+        async def run_list():
+            return [i async for i in f()]
+
+        with self.assertRaisesRegex(Exception, 'aaa'):
+            run_async(run_list())
+
+    def test_comp_8(self):
+        async def f():
+            return [i for i in [1, 2, 3]]
+
+        self.assertEqual(
+            run_async(f()),
+            ([], [1, 2, 3]))
+
+    def test_comp_9(self):
+        async def gen():
+            yield 1
+            yield 2
+        async def f():
+            l = [i async for i in gen()]
+            return [i for i in l]
+
+        self.assertEqual(
+            run_async(f()),
+            ([], [1, 2]))
+
+    def test_comp_10(self):
+        async def f():
+            xx = {i for i in [1, 2, 3]}
+            return {x: x for x in xx}
+
+        self.assertEqual(
+            run_async(f()),
+            ([], {1: 1, 2: 2, 3: 3}))
+
     def test_copy(self):
         async def func(): pass
         coro = func()

Misc/NEWS

@@ -108,7 +108,7 @@ Core and Builtins
   In a brand new thread, raise a RuntimeError since there is no active
   exception to reraise. Patch written by Xiang Zhang.
 
-
+- Issue #28008: Implement PEP 530 -- asynchronous comprehensions.
 
 Library
 -------

Parser/Python.asdl

@@ -110,7 +110,7 @@ module Python
 
     cmpop = Eq | NotEq | Lt | LtE | Gt | GtE | Is | IsNot | In | NotIn
 
-    comprehension = (expr target, expr iter, expr* ifs)
+    comprehension = (expr target, expr iter, expr* ifs, int is_async)
 
     excepthandler = ExceptHandler(expr? type, identifier? name, stmt* body)
                     attributes (int lineno, int col_offset)

Python/Python-ast.c

@@ -435,10 +435,12 @@ static PyTypeObject *NotIn_type;
 static PyTypeObject *comprehension_type;
 static PyObject* ast2obj_comprehension(void*);
 _Py_IDENTIFIER(ifs);
+_Py_IDENTIFIER(is_async);
 static char *comprehension_fields[]={
     "target",
     "iter",
     "ifs",
+    "is_async",
 };
 static PyTypeObject *excepthandler_type;
 static char *excepthandler_attributes[] = {
@@ -1148,7 +1150,7 @@ static int init_types(void)
     NotIn_singleton = PyType_GenericNew(NotIn_type, NULL, NULL);
     if (!NotIn_singleton) return 0;
     comprehension_type = make_type("comprehension", &AST_type,
-                                   comprehension_fields, 3);
+                                   comprehension_fields, 4);
     if (!comprehension_type) return 0;
     if (!add_attributes(comprehension_type, NULL, 0)) return 0;
     excepthandler_type = make_type("excepthandler", &AST_type, NULL, 0);
@@ -2445,7 +2447,8 @@ Index(expr_ty value, PyArena *arena)
 }
 
 comprehension_ty
-comprehension(expr_ty target, expr_ty iter, asdl_seq * ifs, PyArena *arena)
+comprehension(expr_ty target, expr_ty iter, asdl_seq * ifs, int is_async,
+              PyArena *arena)
 {
     comprehension_ty p;
     if (!target) {
@@ -2464,6 +2467,7 @@ comprehension(expr_ty target, expr_ty iter, asdl_seq * ifs, PyArena *arena)
     p->target = target;
     p->iter = iter;
     p->ifs = ifs;
+    p->is_async = is_async;
     return p;
 }
 
@@ -3722,6 +3726,11 @@ ast2obj_comprehension(void* _o)
     if (_PyObject_SetAttrId(result, &PyId_ifs, value) == -1)
         goto failed;
     Py_DECREF(value);
+    value = ast2obj_int(o->is_async);
+    if (!value) goto failed;
+    if (_PyObject_SetAttrId(result, &PyId_is_async, value) == -1)
+        goto failed;
+    Py_DECREF(value);
     return result;
 failed:
     Py_XDECREF(value);
@@ -7146,6 +7155,7 @@ obj2ast_comprehension(PyObject* obj, comprehension_ty* out, PyArena* arena)
     expr_ty target;
     expr_ty iter;
     asdl_seq* ifs;
+    int is_async;
 
     if (_PyObject_HasAttrId(obj, &PyId_target)) {
         int res;
@@ -7193,7 +7203,18 @@ obj2ast_comprehension(PyObject* obj, comprehension_ty* out, PyArena* arena)
         PyErr_SetString(PyExc_TypeError, "required field \"ifs\" missing from comprehension");
         return 1;
     }
-    *out = comprehension(target, iter, ifs, arena);
+    if (_PyObject_HasAttrId(obj, &PyId_is_async)) {
+        int res;
+        tmp = _PyObject_GetAttrId(obj, &PyId_is_async);
+        if (tmp == NULL) goto failed;
+        res = obj2ast_int(tmp, &is_async, arena);
+        if (res != 0) goto failed;
+        Py_CLEAR(tmp);
+    } else {
+        PyErr_SetString(PyExc_TypeError, "required field \"is_async\" missing from comprehension");
+        return 1;
+    }
+    *out = comprehension(target, iter, ifs, is_async, arena);
     return 0;
 failed:
     Py_XDECREF(tmp);

Python/ast.c

@@ -1747,14 +1747,21 @@ static int
 count_comp_fors(struct compiling *c, const node *n)
 {
     int n_fors = 0;
+    int is_async;
 
   count_comp_for:
+    is_async = 0;
     n_fors++;
     REQ(n, comp_for);
-    if (NCH(n) == 5)
-        n = CHILD(n, 4);
-    else
+    if (TYPE(CHILD(n, 0)) == ASYNC) {
+        is_async = 1;
+    }
+    if (NCH(n) == (5 + is_async)) {
+        n = CHILD(n, 4 + is_async);
+    }
+    else {
         return n_fors;
+    }
   count_comp_iter:
     REQ(n, comp_iter);
     n = CHILD(n, 0);
@@ -1817,14 +1824,19 @@ ast_for_comprehension(struct compiling *c, const node *n)
         asdl_seq *t;
         expr_ty expression, first;
         node *for_ch;
+        int is_async = 0;
 
         REQ(n, comp_for);
 
-        for_ch = CHILD(n, 1);
+        if (TYPE(CHILD(n, 0)) == ASYNC) {
+            is_async = 1;
+        }
+
+        for_ch = CHILD(n, 1 + is_async);
         t = ast_for_exprlist(c, for_ch, Store);
         if (!t)
             return NULL;
-        expression = ast_for_expr(c, CHILD(n, 3));
+        expression = ast_for_expr(c, CHILD(n, 3 + is_async));
         if (!expression)
             return NULL;
 
@@ -1832,19 +1844,20 @@ ast_for_comprehension(struct compiling *c, const node *n)
            (x for x, in ...) has 1 element in t, but still requires a Tuple. */
         first = (expr_ty)asdl_seq_GET(t, 0);
         if (NCH(for_ch) == 1)
-            comp = comprehension(first, expression, NULL, c->c_arena);
+            comp = comprehension(first, expression, NULL,
+                                 is_async, c->c_arena);
         else
-            comp = comprehension(Tuple(t, Store, first->lineno, first->col_offset,
-                                     c->c_arena),
-                               expression, NULL, c->c_arena);
+            comp = comprehension(Tuple(t, Store, first->lineno,
+                                       first->col_offset, c->c_arena),
+                                 expression, NULL, is_async, c->c_arena);
         if (!comp)
             return NULL;
 
-        if (NCH(n) == 5) {
+        if (NCH(n) == (5 + is_async)) {
             int j, n_ifs;
             asdl_seq *ifs;
 
-            n = CHILD(n, 4);
+            n = CHILD(n, 4 + is_async);
             n_ifs = count_comp_ifs(c, n);
             if (n_ifs == -1)
                 return NULL;

Python/compile.c

@@ -202,6 +202,16 @@ static int compiler_call_helper(struct compiler *c, int n,
 static int compiler_try_except(struct compiler *, stmt_ty);
 static int compiler_set_qualname(struct compiler *);
 
+static int compiler_sync_comprehension_generator(
+                                      struct compiler *c,
+                                      asdl_seq *generators, int gen_index,
+                                      expr_ty elt, expr_ty val, int type);
+
+static int compiler_async_comprehension_generator(
+                                      struct compiler *c,
+                                      asdl_seq *generators, int gen_index,
+                                      expr_ty elt, expr_ty val, int type);
+
 static PyCodeObject *assemble(struct compiler *, int addNone);
 static PyObject *__doc__;
 
@@ -2165,14 +2175,14 @@ compiler_for(struct compiler *c, stmt_ty s)
 static int
 compiler_async_for(struct compiler *c, stmt_ty s)
 {
-    static PyObject *stopiter_error = NULL;
+    _Py_IDENTIFIER(StopAsyncIteration);
+
     basicblock *try, *except, *end, *after_try, *try_cleanup,
                *after_loop, *after_loop_else;
 
-    if (stopiter_error == NULL) {
-        stopiter_error = PyUnicode_InternFromString("StopAsyncIteration");
-        if (stopiter_error == NULL)
-            return 0;
+    PyObject *stop_aiter_error = _PyUnicode_FromId(&PyId_StopAsyncIteration);
+    if (stop_aiter_error == NULL) {
+        return 0;
     }
 
     try = compiler_new_block(c);
@@ -2214,7 +2224,7 @@ compiler_async_for(struct compiler *c, stmt_ty s)
 
     compiler_use_next_block(c, except);
     ADDOP(c, DUP_TOP);
-    ADDOP_O(c, LOAD_GLOBAL, stopiter_error, names);
+    ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
     ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
     ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
 
@@ -3627,10 +3637,27 @@ compiler_call_helper(struct compiler *c,
     - iterate over the generator sequence instead of using recursion
 */
 
+
 static int
 compiler_comprehension_generator(struct compiler *c,
                                  asdl_seq *generators, int gen_index,
                                  expr_ty elt, expr_ty val, int type)
+{
+    comprehension_ty gen;
+    gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
+    if (gen->is_async) {
+        return compiler_async_comprehension_generator(
+            c, generators, gen_index, elt, val, type);
+    } else {
+        return compiler_sync_comprehension_generator(
+            c, generators, gen_index, elt, val, type);
+    }
+}
+
+static int
+compiler_sync_comprehension_generator(struct compiler *c,
+                                      asdl_seq *generators, int gen_index,
+                                      expr_ty elt, expr_ty val, int type)
 {
     /* generate code for the iterator, then each of the ifs,
        and then write to the element */
@@ -3717,21 +3744,168 @@ compiler_comprehension_generator(struct compiler *c,
     return 1;
 }
 
+static int
+compiler_async_comprehension_generator(struct compiler *c,
+                                      asdl_seq *generators, int gen_index,
+                                      expr_ty elt, expr_ty val, int type)
+{
+    _Py_IDENTIFIER(StopAsyncIteration);
+
+    comprehension_ty gen;
+    basicblock *anchor, *skip, *if_cleanup, *try,
+               *after_try, *except, *try_cleanup;
+    Py_ssize_t i, n;
+
+    PyObject *stop_aiter_error = _PyUnicode_FromId(&PyId_StopAsyncIteration);
+    if (stop_aiter_error == NULL) {
+        return 0;
+    }
+
+    try = compiler_new_block(c);
+    after_try = compiler_new_block(c);
+    try_cleanup = compiler_new_block(c);
+    except = compiler_new_block(c);
+    skip = compiler_new_block(c);
+    if_cleanup = compiler_new_block(c);
+    anchor = compiler_new_block(c);
+
+    if (skip == NULL || if_cleanup == NULL || anchor == NULL ||
+            try == NULL || after_try == NULL ||
+            except == NULL || after_try == NULL) {
+        return 0;
+    }
+
+    gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
+
+    if (gen_index == 0) {
+        /* Receive outermost iter as an implicit argument */
+        c->u->u_argcount = 1;
+        ADDOP_I(c, LOAD_FAST, 0);
+    }
+    else {
+        /* Sub-iter - calculate on the fly */
+        VISIT(c, expr, gen->iter);
+        ADDOP(c, GET_AITER);
+        ADDOP_O(c, LOAD_CONST, Py_None, consts);
+        ADDOP(c, YIELD_FROM);
+    }
+
+    compiler_use_next_block(c, try);
+
+
+    ADDOP_JREL(c, SETUP_EXCEPT, except);
+    if (!compiler_push_fblock(c, EXCEPT, try))
+        return 0;
+
+    ADDOP(c, GET_ANEXT);
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    ADDOP(c, YIELD_FROM);
+    VISIT(c, expr, gen->target);
+    ADDOP(c, POP_BLOCK);
+    compiler_pop_fblock(c, EXCEPT, try);
+    ADDOP_JREL(c, JUMP_FORWARD, after_try);
+
+
+    compiler_use_next_block(c, except);
+    ADDOP(c, DUP_TOP);
+    ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
+    ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
+    ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
+
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
+    ADDOP_JABS(c, JUMP_ABSOLUTE, anchor);
+
+
+    compiler_use_next_block(c, try_cleanup);
+    ADDOP(c, END_FINALLY);
+
+    compiler_use_next_block(c, after_try);
+
+    n = asdl_seq_LEN(gen->ifs);
+    for (i = 0; i < n; i++) {
+        expr_ty e = (expr_ty)asdl_seq_GET(gen->ifs, i);
+        VISIT(c, expr, e);
+        ADDOP_JABS(c, POP_JUMP_IF_FALSE, if_cleanup);
+        NEXT_BLOCK(c);
+    }
+
+    if (++gen_index < asdl_seq_LEN(generators))
+        if (!compiler_comprehension_generator(c,
+                                              generators, gen_index,
+                                              elt, val, type))
+        return 0;
+
+    /* only append after the last for generator */
+    if (gen_index >= asdl_seq_LEN(generators)) {
+        /* comprehension specific code */
+        switch (type) {
+        case COMP_GENEXP:
+            VISIT(c, expr, elt);
+            ADDOP(c, YIELD_VALUE);
+            ADDOP(c, POP_TOP);
+            break;
+        case COMP_LISTCOMP:
+            VISIT(c, expr, elt);
+            ADDOP_I(c, LIST_APPEND, gen_index + 1);
+            break;
+        case COMP_SETCOMP:
+            VISIT(c, expr, elt);
+            ADDOP_I(c, SET_ADD, gen_index + 1);
+            break;
+        case COMP_DICTCOMP:
+            /* With 'd[k] = v', v is evaluated before k, so we do
+               the same. */
+            VISIT(c, expr, val);
+            VISIT(c, expr, elt);
+            ADDOP_I(c, MAP_ADD, gen_index + 1);
+            break;
+        default:
+            return 0;
+        }
+
+        compiler_use_next_block(c, skip);
+    }
+    compiler_use_next_block(c, if_cleanup);
+    ADDOP_JABS(c, JUMP_ABSOLUTE, try);
+    compiler_use_next_block(c, anchor);
+    ADDOP(c, POP_TOP);
+
+    return 1;
+}
+
 static int
 compiler_comprehension(struct compiler *c, expr_ty e, int type,
                        identifier name, asdl_seq *generators, expr_ty elt,
                        expr_ty val)
 {
     PyCodeObject *co = NULL;
-    expr_ty outermost_iter;
+    comprehension_ty outermost;
     PyObject *qualname = NULL;
+    int is_async_function = c->u->u_ste->ste_coroutine;
+    int is_async_generator = 0;
 
-    outermost_iter = ((comprehension_ty)
-                      asdl_seq_GET(generators, 0))->iter;
+    outermost = (comprehension_ty) asdl_seq_GET(generators, 0);
 
     if (!compiler_enter_scope(c, name, COMPILER_SCOPE_COMPREHENSION,
                               (void *)e, e->lineno))
+    {
         goto error;
+    }
+
+    is_async_generator = c->u->u_ste->ste_coroutine;
+
+    if (is_async_generator && !is_async_function) {
+        if (e->lineno > c->u->u_lineno) {
+            c->u->u_lineno = e->lineno;
+            c->u->u_lineno_set = 0;
+        }
+        compiler_error(c, "asynchronous comprehension outside of "
+                          "an asynchronous function");
+        goto error_in_scope;
+    }
 
     if (type != COMP_GENEXP) {
         int op;
@@ -3774,9 +3948,24 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
     Py_DECREF(qualname);
     Py_DECREF(co);
 
-    VISIT(c, expr, outermost_iter);
-    ADDOP(c, GET_ITER);
+    VISIT(c, expr, outermost->iter);
+
+    if (outermost->is_async) {
+        ADDOP(c, GET_AITER);
+        ADDOP_O(c, LOAD_CONST, Py_None, consts);
+        ADDOP(c, YIELD_FROM);
+    } else {
+        ADDOP(c, GET_ITER);
+    }
+
     ADDOP_I(c, CALL_FUNCTION, 1);
+
+    if (is_async_generator && type != COMP_GENEXP) {
+        ADDOP(c, GET_AWAITABLE);
+        ADDOP_O(c, LOAD_CONST, Py_None, consts);
+        ADDOP(c, YIELD_FROM);
+    }
+
     return 1;
 error_in_scope:
     compiler_exit_scope(c);
@@ -4140,11 +4329,8 @@ compiler_visit_expr(struct compiler *c, expr_ty e)
         if (c->u->u_ste->ste_type != FunctionBlock)
             return compiler_error(c, "'await' outside function");
 
-        if (c->u->u_scope_type == COMPILER_SCOPE_COMPREHENSION)
-            return compiler_error(
-                c, "'await' expressions in comprehensions are not supported");
-
-        if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION)
+        if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION &&
+                c->u->u_scope_type != COMPILER_SCOPE_COMPREHENSION)
             return compiler_error(c, "'await' outside async function");
 
         VISIT(c, expr, e->v.Await.value);

Python/graminit.c

@@ -1812,32 +1812,37 @@ static state states_80[2] = {
     {2, arcs_80_0},
     {1, arcs_80_1},
 };
-static arc arcs_81_0[1] = {
-    {101, 1},
+static arc arcs_81_0[2] = {
+    {21, 1},
+    {101, 2},
 };
 static arc arcs_81_1[1] = {
-    {66, 2},
+    {101, 2},
 };
 static arc arcs_81_2[1] = {
-    {102, 3},
+    {66, 3},
 };
 static arc arcs_81_3[1] = {
-    {112, 4},
+    {102, 4},
 };
-static arc arcs_81_4[2] = {
-    {171, 5},
-    {0, 4},
+static arc arcs_81_4[1] = {
+    {112, 5},
 };
-static arc arcs_81_5[1] = {
+static arc arcs_81_5[2] = {
+    {171, 6},
     {0, 5},
 };
-static state states_81[6] = {
-    {1, arcs_81_0},
+static arc arcs_81_6[1] = {
+    {0, 6},
+};
+static state states_81[7] = {
+    {2, arcs_81_0},
     {1, arcs_81_1},
     {1, arcs_81_2},
     {1, arcs_81_3},
-    {2, arcs_81_4},
-    {1, arcs_81_5},
+    {1, arcs_81_4},
+    {2, arcs_81_5},
+    {1, arcs_81_6},
 };
 static arc arcs_82_0[1] = {
     {97, 1},
@@ -2060,9 +2065,9 @@ static dfa dfas[86] = {
     {335, "argument", 0, 4, states_79,
      "\000\040\200\000\006\000\000\000\000\000\010\000\000\000\020\002\000\300\220\050\037\000"},
     {336, "comp_iter", 0, 2, states_80,
-     "\000\000\000\000\000\000\000\000\000\000\000\000\042\000\000\000\000\000\000\000\000\000"},
-    {337, "comp_for", 0, 6, states_81,
-     "\000\000\000\000\000\000\000\000\000\000\000\000\040\000\000\000\000\000\000\000\000\000"},
+     "\000\000\040\000\000\000\000\000\000\000\000\000\042\000\000\000\000\000\000\000\000\000"},
+    {337, "comp_for", 0, 7, states_81,
+     "\000\000\040\000\000\000\000\000\000\000\000\000\040\000\000\000\000\000\000\000\000\000"},
     {338, "comp_if", 0, 4, states_82,
      "\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000"},
     {339, "encoding_decl", 0, 2, states_83,

Python/symtable.c

@@ -1682,6 +1682,9 @@ symtable_visit_comprehension(struct symtable *st, comprehension_ty lc)
     VISIT(st, expr, lc->target);
     VISIT(st, expr, lc->iter);
     VISIT_SEQ(st, expr, lc->ifs);
+    if (lc->is_async) {
+        st->st_cur->ste_coroutine = 1;
+    }
     return 1;
 }
 
@@ -1734,6 +1737,9 @@ symtable_handle_comprehension(struct symtable *st, expr_ty e,
         return 0;
     }
     st->st_cur->ste_generator = is_generator;
+    if (outermost->is_async) {
+        st->st_cur->ste_coroutine = 1;
+    }
     /* Outermost iter is received as an argument */
     if (!symtable_implicit_arg(st, 0)) {
         symtable_exit_block(st, (void *)e);