From e2d349f781c5f21efe58c9dba7a221f65bcbdb55 Mon Sep 17 00:00:00 2001
From: rocky <rb@dustyfeet.com>
Date: Tue, 31 Mar 2020 12:05:39 -0400
Subject: [PATCH] Handle nested async for in for...

and Better async comprehension detection.

Still more work is needed. See commented-out section in
test/simple_source/bug37/02_async_for_generator.py
---
 test/bytecode_3.7/02_async_for_generator.pyc  | Bin 692 -> 991 bytes
 .../bug37/02_async_for_generator.py           |   8 ++
 uncompyle6/parsers/parse37base.py             | 134 +++++++++++++-----
 uncompyle6/semantics/customize3.py            |  21 ++-
 uncompyle6/semantics/customize35.py           |  11 +-
 uncompyle6/semantics/pysource.py              |   8 +-
 6 files changed, 130 insertions(+), 52 deletions(-)

diff --git a/test/bytecode_3.7/02_async_for_generator.pyc b/test/bytecode_3.7/02_async_for_generator.pyc
index cc9c8b33893eea37eaef355c6b109eb662815f84..935042aaeabf2e68638acf487c1795e81127eea9 100644
GIT binary patch
delta 376
zcmXv}%Syvg5S_WXX=6-EmtAz_LWrP<Zd^!_;zBpsxbOigC6`Jt_Hh#zN=p|h*lj8J
z0sevD*JR_bcqZUn&Y8K7b7t<3``c{w!_X(hS6V)N^19X8>0)Qd@RIJxDB~cr(Z)qq
zrg244^ZTr)7KsuWLb3ytK;8j;(2xxqwe$))r;;yWcR>(3al*Kl5WzOE0rC^j6V2%U
zaLy|_qPO6<mdw>PjeJc7J^T6I4b9ZrnFdU<7Pp|OhAYPrTNxs9tk_n1GDxk}a>5l4
zt&7Swgyg3IkjwK1y?5>F=rK$4;@~tJCxa~Sp};m7zf5`}o<z!qqHGC?5qj12z8F^C
ezk&Rvb|toL)G>zHLblS*wl^!a>&Zl(_VgcE*g$Xq

delta 75
zcmcc5zJ-;~iI<m)0SJ@@zsEjin8+s~B?9E7Fr+ZHFhnt=Fr_d9X~q<mU<OUrjUIMP
YTz;BNMVvq-MLd%?GOMvM0ht^u07KLb)&Kwi

diff --git a/test/simple_source/bug37/02_async_for_generator.py b/test/simple_source/bug37/02_async_for_generator.py
index 1a497b3f..54c40a98 100644
--- a/test/simple_source/bug37/02_async_for_generator.py
+++ b/test/simple_source/bug37/02_async_for_generator.py
@@ -6,3 +6,11 @@ def make_arange(n):
 
 async def run(m):
     return [i async for i in m]
+
+# From 3.7.6 test_coroutines.py
+async def run_list(pair, f):
+    return [i for pair in p async for i in f]
+
+# FIXME: add this. It works in decompyle3
+# async def run_gen():
+#     return (i async for i in f if 0 < i < 4)
diff --git a/uncompyle6/parsers/parse37base.py b/uncompyle6/parsers/parse37base.py
index 2cd9a2b4..3e58ef6c 100644
--- a/uncompyle6/parsers/parse37base.py
+++ b/uncompyle6/parsers/parse37base.py
@@ -210,41 +210,86 @@ class Python37BaseParser(PythonParser):
 
                 if self.version < 3.8:
                     rules_str += """
-                      stmt                ::= async_with_stmt SETUP_ASYNC_WITH
-                      async_with_stmt     ::= expr
+                      stmt                 ::= async_with_stmt SETUP_ASYNC_WITH
+                      c_stmt               ::= c_async_with_stmt SETUP_ASYNC_WITH
+                      async_with_stmt      ::= expr
+                                               async_with_pre
+                                               POP_TOP
+                                               suite_stmts_opt
+                                               POP_BLOCK LOAD_CONST
+                                               async_with_post
+                      c_async_with_stmt    ::= expr
+                                               async_with_pre
+                                               POP_TOP
+                                               c_suite_stmts_opt
+                                               POP_BLOCK LOAD_CONST
+                                               async_with_post
+                      async_with_stmt      ::= expr
+                                               async_with_pre
+                                               POP_TOP
+                                               suite_stmts_opt
+                                               async_with_post
+                      c_async_with_stmt    ::= expr
+                                               async_with_pre
+                                               POP_TOP
+                                               c_suite_stmts_opt
+                                               async_with_post
+                      async_with_as_stmt   ::= expr
+                                               async_with_pre
+                                               store
+                                               suite_stmts_opt
+                                               POP_BLOCK LOAD_CONST
+                                               async_with_post
+                      c_async_with_as_stmt ::= expr
                                               async_with_pre
-                                              POP_TOP
-                                              suite_stmts_opt
+                                              store
+                                              c_suite_stmts_opt
                                               POP_BLOCK LOAD_CONST
                                               async_with_post
-                      async_with_stmt     ::= expr
-                                              async_with_pre
-                                              POP_TOP
-                                              suite_stmts_opt
-                                              async_with_post
-                      async_with_as_stmt  ::= expr
+                      async_with_as_stmt   ::= expr
+                                              async_with_pre
+                                              store
+                                              suite_stmts_opt
+                                              async_with_post
+                      c_async_with_as_stmt ::= expr
                                               async_with_pre
                                               store
                                               suite_stmts_opt
-                                              POP_BLOCK LOAD_CONST
                                               async_with_post
                     """
                 else:
                     rules_str += """
-                      async_with_pre      ::= BEFORE_ASYNC_WITH GET_AWAITABLE LOAD_CONST YIELD_FROM SETUP_ASYNC_WITH
-                      async_with_post     ::= BEGIN_FINALLY COME_FROM_ASYNC_WITH
-                                              WITH_CLEANUP_START GET_AWAITABLE LOAD_CONST YIELD_FROM
-                                              WITH_CLEANUP_FINISH END_FINALLY
-                      async_with_stmt     ::= expr
+                      async_with_pre       ::= BEFORE_ASYNC_WITH GET_AWAITABLE LOAD_CONST YIELD_FROM SETUP_ASYNC_WITH
+                      async_with_post      ::= BEGIN_FINALLY COME_FROM_ASYNC_WITH
+                                               WITH_CLEANUP_START GET_AWAITABLE LOAD_CONST YIELD_FROM
+                                               WITH_CLEANUP_FINISH END_FINALLY
+                      async_with_stmt      ::= expr
+                                               async_with_pre
+                                               POP_TOP
+                                               suite_stmts
+                                               POP_TOP POP_BLOCK
+                                               async_with_post
+                      c_async_with_stmt    ::= expr
+                                               async_with_pre
+                                               POP_TOP
+                                               c_suite_stmts
+                                               POP_TOP POP_BLOCK
+                                               async_with_post
+                      async_with_stmt      ::= expr
+                                               async_with_pre
+                                               POP_TOP
+                                               suite_stmts
+                                               POP_BLOCK
+                                               BEGIN_FINALLY
+                                               WITH_CLEANUP_START GET_AWAITABLE LOAD_CONST YIELD_FROM
+                                               WITH_CLEANUP_FINISH POP_FINALLY LOAD_CONST RETURN_VALUE
+                                               COME_FROM_ASYNC_WITH
+                                               WITH_CLEANUP_START GET_AWAITABLE LOAD_CONST YIELD_FROM
+                                               WITH_CLEANUP_FINISH END_FINALLY
+                      c_async_with_stmt   ::= expr
                                               async_with_pre
                                               POP_TOP
-                                              suite_stmts
-                                              POP_TOP POP_BLOCK
-                                              async_with_post
-                      async_with_stmt     ::= expr
-                                              async_with_pre
-                                              POP_TOP
-                                              suite_stmts
+                                              c_suite_stmts
                                               POP_BLOCK
                                               BEGIN_FINALLY
                                               WITH_CLEANUP_START GET_AWAITABLE LOAD_CONST YIELD_FROM
@@ -252,15 +297,24 @@ class Python37BaseParser(PythonParser):
                                               COME_FROM_ASYNC_WITH
                                               WITH_CLEANUP_START GET_AWAITABLE LOAD_CONST YIELD_FROM
                                               WITH_CLEANUP_FINISH END_FINALLY
-                      async_with_as_stmt  ::= expr
-                                              async_with_pre
-                                              store suite_stmts
-                                              POP_TOP POP_BLOCK
-                                              async_with_post
-                      async_with_as_stmt  ::= expr
-                                              async_with_pre
-                                              store suite_stmts
-                                              POP_BLOCK async_with_post
+                      async_with_as_stmt   ::= expr
+                                               async_with_pre
+                                               store suite_stmts
+                                               POP_TOP POP_BLOCK
+                                               async_with_post
+                      c_async_with_as_stmt ::= expr
+                                               async_with_pre
+                                               store suite_stmts
+                                               POP_TOP POP_BLOCK
+                                               async_with_post
+                      async_with_as_stmt   ::= expr
+                                               async_with_pre
+                                               store suite_stmts
+                                               POP_BLOCK async_with_post
+                      c_async_with_as_stmt ::= expr
+                                               async_with_pre
+                                               store suite_stmts
+                                               POP_BLOCK async_with_post
                     """
                 self.addRule(rules_str, nop_func)
 
@@ -539,7 +593,7 @@ class Python37BaseParser(PythonParser):
 
                     stmt                ::= genexpr_func_async
 
-                    func_async_prefix   ::= SETUP_EXCEPT GET_ANEXT LOAD_CONST YIELD_FROM
+                    func_async_prefix   ::= _come_froms SETUP_EXCEPT GET_ANEXT LOAD_CONST YIELD_FROM
                     func_async_middle   ::= POP_BLOCK JUMP_FORWARD COME_FROM_EXCEPT
                                             DUP_TOP LOAD_GLOBAL COMPARE_OP POP_JUMP_IF_TRUE
                                             END_FINALLY COME_FROM
@@ -548,22 +602,24 @@ class Python37BaseParser(PythonParser):
                                             JUMP_BACK COME_FROM
                                             POP_TOP POP_TOP POP_TOP POP_EXCEPT POP_TOP
 
-                    expr                ::= listcomp_async
-                    listcomp_async      ::= LOAD_LISTCOMP LOAD_STR MAKE_FUNCTION_0
+                    expr                ::= list_comp_async
+                    list_comp_async     ::= LOAD_LISTCOMP LOAD_STR MAKE_FUNCTION_0
                                             expr GET_AITER CALL_FUNCTION_1
                                             GET_AWAITABLE LOAD_CONST
                                             YIELD_FROM
 
-                    expr                 ::= listcomp_async
-                    listcomp_async       ::= BUILD_LIST_0 LOAD_FAST func_async_prefix
+                    expr                ::= list_comp_async
+                    list_afor2          ::= func_async_prefix
                                             store func_async_middle list_iter
                                             JUMP_BACK COME_FROM
                                             POP_TOP POP_TOP POP_TOP POP_EXCEPT POP_TOP
-
+                    list_comp_async     ::= BUILD_LIST_0 LOAD_FAST list_afor2
+                    get_aiter           ::= LOAD_DEREF GET_AITER
+                    list_afor           ::= get_aiter list_afor2
+                    list_iter           ::= list_afor
                    """,
                     nop_func,
                 )
-                custom_ops_processed.add(opname)
             elif opname == "JUMP_IF_NOT_DEBUG":
                 v = token.attr
                 self.addRule(
diff --git a/uncompyle6/semantics/customize3.py b/uncompyle6/semantics/customize3.py
index f1394fd3..ea7c55f6 100644
--- a/uncompyle6/semantics/customize3.py
+++ b/uncompyle6/semantics/customize3.py
@@ -17,6 +17,7 @@
 """
 
 from uncompyle6.semantics.consts import TABLE_DIRECT
+from xdis.util import co_flags_is_async
 
 from xdis.code import iscode
 from uncompyle6.scanner import Code
@@ -87,7 +88,9 @@ def customize_for_version3(self, version):
         p = self.prec
         self.prec = 27
 
-        code = Code(node[1].attr, self.scanner, self.currentclass)
+        code_obj = node[1].attr
+        assert iscode(code_obj)
+        code = Code(code_obj, self.scanner, self.currentclass)
         ast = self.build_ast(code._tokens, code._customize)
         self.customize(code._customize)
 
@@ -141,7 +144,10 @@ def customize_for_version3(self, version):
             # Find the list comprehension body. It is the inner-most
             # node that is not list_.. .
             while n == "list_iter":
-                n = n[0]  # recurse one step
+
+                # recurse one step
+                n = n[0]
+
                 if n == "list_for":
                     stores.append(n[2])
                     n = n[3]
@@ -168,6 +174,11 @@ def customize_for_version3(self, version):
                     list_ifs.append(n)
                     n = n[-1]
                     pass
+                elif n == "list_afor":
+                    collections.append(n[0][0])
+                    n = n[1]
+                    stores.append(n[1][0])
+                    n = n[3]
                 pass
 
             assert n == "lc_body", ast
@@ -175,7 +186,13 @@ def customize_for_version3(self, version):
             self.preorder(n[0])
 
         # FIXME: add indentation around "for"'s and "in"'s
+        n_colls = len(collections)
         for i, store in enumerate(stores):
+            if i >= n_colls:
+                break
+            if collections[i] == "LOAD_DEREF"  and co_flags_is_async(code_obj.co_flags):
+                self.write(" async")
+                pass
             self.write(" for ")
             self.preorder(store)
             self.write(" in ")
diff --git a/uncompyle6/semantics/customize35.py b/uncompyle6/semantics/customize35.py
index 23c3b6d2..62be6b18 100644
--- a/uncompyle6/semantics/customize35.py
+++ b/uncompyle6/semantics/customize35.py
@@ -16,7 +16,7 @@
 """
 
 from xdis.code import iscode
-from xdis.util import COMPILER_FLAG_BIT
+from xdis.util import co_flags_is_async
 from uncompyle6.semantics.consts import (
     INDENT_PER_LEVEL,
     PRECEDENCE,
@@ -207,14 +207,7 @@ def customize_for_version35(self, version):
         pass
 
         is_code = hasattr(code_node, "attr") and iscode(code_node.attr)
-        return is_code and (
-            code_node.attr.co_flags
-            & (
-                COMPILER_FLAG_BIT["COROUTINE"]
-                | COMPILER_FLAG_BIT["ITERABLE_COROUTINE"]
-                | COMPILER_FLAG_BIT["ASYNC_GENERATOR"]
-            )
-        )
+        return is_code and co_flags_is_async(code_node.attr.co_flags)
 
     def n_function_def(node):
         if is_async_fn(node):
diff --git a/uncompyle6/semantics/pysource.py b/uncompyle6/semantics/pysource.py
index 0347f0f0..4cd63a43 100644
--- a/uncompyle6/semantics/pysource.py
+++ b/uncompyle6/semantics/pysource.py
@@ -1207,7 +1207,11 @@ class SourceWalker(GenericASTTraversal, object):
 
         is_30_dict_comp = False
         store = None
-        n = ast[iter_index]
+        if node == "list_comp_async":
+            n = ast[2][1]
+        else:
+            n = ast[iter_index]
+
         if ast in (
             "set_comp_func",
             "dict_comp_func",
@@ -1238,7 +1242,7 @@ class SourceWalker(GenericASTTraversal, object):
                     pass
                 pass
         elif ast == "listcomp_async":
-            store = ast[3]
+            store = ast[2][1]
         else:
             assert n == "list_iter", n