bpo-33346: Allow async comprehensions inside implicit async comprehensions (GH-6766)

Co-authored-by: Pablo Galindo <[email protected]>

Author: serhiy-storchaka

Doc/reference/expressions.rst

@@ -218,9 +218,9 @@ A comprehension in an :keyword:`!async def` function may consist of either a
 :keyword:`!for` or :keyword:`!async for` clause following the leading
 expression, may contain additional :keyword:`!for` or :keyword:`!async for`
 clauses, and may also use :keyword:`await` expressions.
-If a comprehension contains either :keyword:`!async for` clauses
-or :keyword:`!await` expressions it is called an
-:dfn:`asynchronous comprehension`.  An asynchronous comprehension may
+If a comprehension contains either :keyword:`!async for` clauses or
+:keyword:`!await` expressions or other asynchronous comprehensions it is called
+an :dfn:`asynchronous comprehension`.  An asynchronous comprehension may
 suspend the execution of the coroutine function in which it appears.
 See also :pep:`530`.
 
@@ -230,6 +230,11 @@ See also :pep:`530`.
 .. versionchanged:: 3.8
    ``yield`` and ``yield from`` prohibited in the implicitly nested scope.
 
+.. versionchanged:: 3.11
+   Asynchronous comprehensions are now allowed inside comprehensions in
+   asynchronous functions. Outer comprehensions implicitly become
+   asynchronous.
+
 
 .. _lists:
 

Doc/whatsnew/3.11.rst

@@ -148,16 +148,19 @@ See :pep:`657` for more details. (Contributed by Pablo Galindo, Batuhan Taskaya
 and Ammar Askar in :issue:`43950`.)
 
 
-
 Other Language Changes
 ======================
 
-A :exc:`TypeError` is now raised instead of an :exc:`AttributeError` in
-:meth:`contextlib.ExitStack.enter_context` and
-:meth:`contextlib.AsyncExitStack.enter_async_context` for objects which do not
-support the :term:`context manager` or :term:`asynchronous context manager`
-protocols correspondingly.
-(Contributed by Serhiy Storchaka in :issue:`44471`.)
+* Asynchronous comprehensions are now allowed inside comprehensions in
+  asynchronous functions. Outer comprehensions implicitly become
+  asynchronous. (Contributed by Serhiy Storchaka in :issue:`33346`.)
+
+* A :exc:`TypeError` is now raised instead of an :exc:`AttributeError` in
+  :meth:`contextlib.ExitStack.enter_context` and
+  :meth:`contextlib.AsyncExitStack.enter_async_context` for objects which do not
+  support the :term:`context manager` or :term:`asynchronous context manager`
+  protocols correspondingly.
+  (Contributed by Serhiy Storchaka in :issue:`44471`.)
 
 * A :exc:`TypeError` is now raised instead of an :exc:`AttributeError` in
   :keyword:`with` and :keyword:`async with` statements for objects which do not

Lib/test/test_coroutines.py

@@ -28,6 +28,12 @@ def __await__(self):
         yield self.value
 
 
+async def asynciter(iterable):
+    """Convert an iterable to an asynchronous iterator."""
+    for x in iterable:
+        yield x
+
+
 def run_async(coro):
     assert coro.__class__ in {types.GeneratorType, types.CoroutineType}
 
@@ -125,6 +131,11 @@ def bar():
                     for c in b]
             """,
 
+            """async def foo():
+                def bar():
+                 [[async for i in b] for b in els]
+            """,
+
             """async def foo():
                 def bar():
                  [i for i in els
@@ -200,6 +211,13 @@ def bar():
                  [i for i in els if await i]
             """,
 
+            """def bar():
+                 [[i async for i in a] for a in elts]
+            """,
+
+            """[[i async for i in a] for a in elts]
+            """,
+
             """async def foo():
                 await
             """,
@@ -2011,6 +2029,60 @@ async def f():
             run_async(f()),
             ([], {1: 1, 2: 2, 3: 3}))
 
+    def test_nested_comp(self):
+        async def run_list_inside_list():
+            return [[i + j async for i in asynciter([1, 2])] for j in [10, 20]]
+        self.assertEqual(
+            run_async(run_list_inside_list()),
+            ([], [[11, 12], [21, 22]]))
+
+        async def run_set_inside_list():
+            return [{i + j async for i in asynciter([1, 2])} for j in [10, 20]]
+        self.assertEqual(
+            run_async(run_set_inside_list()),
+            ([], [{11, 12}, {21, 22}]))
+
+        async def run_list_inside_set():
+            return {sum([i async for i in asynciter(range(j))]) for j in [3, 5]}
+        self.assertEqual(
+            run_async(run_list_inside_set()),
+            ([], {3, 10}))
+
+        async def run_dict_inside_dict():
+            return {j: {i: i + j async for i in asynciter([1, 2])} for j in [10, 20]}
+        self.assertEqual(
+            run_async(run_dict_inside_dict()),
+            ([], {10: {1: 11, 2: 12}, 20: {1: 21, 2: 22}}))
+
+        async def run_list_inside_gen():
+            gen = ([i + j async for i in asynciter([1, 2])] for j in [10, 20])
+            return [x async for x in gen]
+        self.assertEqual(
+            run_async(run_list_inside_gen()),
+            ([], [[11, 12], [21, 22]]))
+
+        async def run_gen_inside_list():
+            gens = [(i async for i in asynciter(range(j))) for j in [3, 5]]
+            return [x for g in gens async for x in g]
+        self.assertEqual(
+            run_async(run_gen_inside_list()),
+            ([], [0, 1, 2, 0, 1, 2, 3, 4]))
+
+        async def run_gen_inside_gen():
+            gens = ((i async for i in asynciter(range(j))) for j in [3, 5])
+            return [x for g in gens async for x in g]
+        self.assertEqual(
+            run_async(run_gen_inside_gen()),
+            ([], [0, 1, 2, 0, 1, 2, 3, 4]))
+
+        async def run_list_inside_list_inside_list():
+            return [[[i + j + k async for i in asynciter([1, 2])]
+                     for j in [10, 20]]
+                    for k in [100, 200]]
+        self.assertEqual(
+            run_async(run_list_inside_list_inside_list()),
+            ([], [[[111, 112], [121, 122]], [[211, 212], [221, 222]]]))
+
     def test_copy(self):
         async def func(): pass
         coro = func()

Misc/NEWS.d/next/Core and Builtins/2018-05-11-12-44-03.bpo-33346.ZgBkvB.rst

@@ -0,0 +1,3 @@
+Asynchronous comprehensions are now allowed inside comprehensions in
+asynchronous functions.  Outer comprehensions implicitly become
+asynchronous.

Python/compile.c

@@ -4947,11 +4947,9 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
     PyCodeObject *co = NULL;
     comprehension_ty outermost;
     PyObject *qualname = NULL;
+    int scope_type = c->u->u_scope_type;
     int is_async_generator = 0;
-    int top_level_await = IS_TOP_LEVEL_AWAIT(c);
-
-
-    int is_async_function = c->u->u_ste->ste_coroutine;
+    int is_top_level_await = IS_TOP_LEVEL_AWAIT(c);
 
     outermost = (comprehension_ty) asdl_seq_GET(generators, 0);
     if (!compiler_enter_scope(c, name, COMPILER_SCOPE_COMPREHENSION,
@@ -4963,7 +4961,11 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
 
     is_async_generator = c->u->u_ste->ste_coroutine;
 
-    if (is_async_generator && !is_async_function && type != COMP_GENEXP  && !top_level_await) {
+    if (is_async_generator && type != COMP_GENEXP &&
+        scope_type != COMPILER_SCOPE_ASYNC_FUNCTION &&
+        scope_type != COMPILER_SCOPE_COMPREHENSION &&
+        !is_top_level_await)
+    {
         compiler_error(c, "asynchronous comprehension outside of "
                           "an asynchronous function");
         goto error_in_scope;
@@ -5002,7 +5004,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
     qualname = c->u->u_qualname;
     Py_INCREF(qualname);
     compiler_exit_scope(c);
-    if (top_level_await && is_async_generator){
+    if (is_top_level_await && is_async_generator){
         c->u->u_ste->ste_coroutine = 1;
     }
     if (co == NULL)

Python/symtable.c

@@ -2056,7 +2056,14 @@ symtable_handle_comprehension(struct symtable *st, expr_ty e,
         return 0;
     }
     st->st_cur->ste_generator = is_generator;
-    return symtable_exit_block(st);
+    int is_async = st->st_cur->ste_coroutine && !is_generator;
+    if (!symtable_exit_block(st)) {
+        return 0;
+    }
+    if (is_async) {
+        st->st_cur->ste_coroutine = 1;
+    }
+    return 1;
 }
 
 static int