Use LOAD_ARG in 3.6

uncompyle6/parsers/parse3.py

@@ -65,7 +65,9 @@ class Python3Parser(PythonParser):
 
         list_comp ::= BUILD_LIST_0 list_iter
         lc_body   ::= expr LIST_APPEND
-        list_for ::= expr FOR_ITER store list_iter jb_or_c
+        list_for  ::= expr_or_arg
+                      FOR_ITER
+                      store list_iter jb_or_c
 
         # This is seen in PyPy, but possibly it appears on other Python 3?
         list_if     ::= expr jmp_false list_iter COME_FROM
@@ -77,10 +79,10 @@ class Python3Parser(PythonParser):
 
         stmt ::= set_comp_func
 
-        set_comp_func ::= BUILD_SET_0 LOAD_FAST FOR_ITER store comp_iter
+        set_comp_func ::= BUILD_SET_0 LOAD_ARG FOR_ITER store comp_iter
                           JUMP_BACK RETURN_VALUE RETURN_LAST
 
-        set_comp_func ::= BUILD_SET_0 LOAD_FAST FOR_ITER store comp_iter
+        set_comp_func ::= BUILD_SET_0 LOAD_ARG FOR_ITER store comp_iter
                           COME_FROM JUMP_BACK RETURN_VALUE RETURN_LAST
 
         comp_body ::= dict_comp_body
@@ -88,6 +90,8 @@ class Python3Parser(PythonParser):
         dict_comp_body ::= expr expr MAP_ADD
         set_comp_body ::= expr SET_ADD
 
+        expr_or_arg     ::= LOAD_ARG
+        expr_or_arg     ::= expr
         # See also common Python p_list_comprehension
         """
 
@@ -95,7 +99,7 @@ class Python3Parser(PythonParser):
         """"
         expr ::= dict_comp
         stmt ::= dict_comp_func
-        dict_comp_func ::= BUILD_MAP_0 LOAD_FAST FOR_ITER store
+        dict_comp_func ::= BUILD_MAP_0 LOAD_ARG FOR_ITER store
                            comp_iter JUMP_BACK RETURN_VALUE RETURN_LAST
 
         comp_iter     ::= comp_if_not

uncompyle6/parsers/parse36.py

@@ -354,7 +354,8 @@ class Python36Parser(Python35Parser):
 
                     expr                ::= list_comp_async
                     list_comp_async     ::= LOAD_LISTCOMP LOAD_STR MAKE_FUNCTION_0
-                                            expr GET_AITER CALL_FUNCTION_1
+                                            expr GET_AITER
+                                            LOAD_CONST YIELD_FROM CALL_FUNCTION_1
                                             GET_AWAITABLE LOAD_CONST
                                             YIELD_FROM
 
@@ -370,6 +371,114 @@ class Python36Parser(Python35Parser):
                    """,
                     nop_func,
                 )
+
+            elif opname == "GET_AITER":
+                self.add_unique_doc_rules("get_aiter ::= expr GET_AITER", customize)
+
+                if not {"MAKE_FUNCTION_0", "MAKE_FUNCTION_CLOSURE"} in self.seen_ops:
+                    self.addRule(
+                        """
+                        expr                ::= dict_comp_async
+                        expr                ::= generator_exp_async
+                        expr                ::= list_comp_async
+
+                        dict_comp_async     ::= LOAD_DICTCOMP
+                                                LOAD_STR
+                                                MAKE_FUNCTION_0
+                                                get_aiter
+                                                CALL_FUNCTION_1
+
+                        dict_comp_async     ::= BUILD_MAP_0 LOAD_ARG
+                                                dict_comp_async
+
+                        func_async_middle   ::= POP_BLOCK JUMP_FORWARD COME_FROM_EXCEPT
+                                                DUP_TOP LOAD_GLOBAL COMPARE_OP POP_JUMP_IF_TRUE
+                                                END_FINALLY COME_FROM
+
+                        func_async_prefix   ::= _come_froms SETUP_EXCEPT GET_ANEXT LOAD_CONST YIELD_FROM
+
+                        generator_exp_async ::= load_genexpr LOAD_STR MAKE_FUNCTION_0
+                                                get_aiter CALL_FUNCTION_1
+
+                        genexpr_func_async  ::= LOAD_ARG func_async_prefix
+                                                store func_async_middle comp_iter
+                                                JUMP_LOOP COME_FROM
+                                                POP_TOP POP_TOP POP_TOP POP_EXCEPT POP_TOP
+
+                        # FIXME this is a workaround for probalby some bug in the Earley parser
+                        # if we use get_aiter, then list_comp_async doesn't match, and I don't
+                        # understand why.
+                        expr_get_aiter      ::= expr GET_AITER
+
+                        list_afor           ::= get_aiter list_afor2
+
+                        list_afor2          ::= func_async_prefix
+                                                store func_async_middle list_iter
+                                                JUMP_LOOP COME_FROM
+                                                POP_TOP POP_TOP POP_TOP POP_EXCEPT POP_TOP
+
+                        list_comp_async     ::= BUILD_LIST_0 LOAD_ARG list_afor2
+                        list_comp_async     ::= LOAD_LISTCOMP LOAD_STR MAKE_FUNCTION_0
+                                                expr_get_aiter CALL_FUNCTION_1
+                                                GET_AWAITABLE LOAD_CONST
+                                                YIELD_FROM
+
+                        list_iter           ::= list_afor
+
+                        set_comp_async       ::= LOAD_SETCOMP
+                                                 LOAD_STR
+                                                 MAKE_FUNCTION_0
+                                                 get_aiter
+                                                 CALL_FUNCTION_1
+
+                        set_comp_async       ::= LOAD_CLOSURE
+                                                 BUILD_TUPLE_1
+                                                 LOAD_SETCOMP
+                                                 LOAD_STR MAKE_FUNCTION_CLOSURE
+                                                 get_aiter CALL_FUNCTION_1
+                                                 await
+                       """,
+                        nop_func,
+                    )
+                    custom_ops_processed.add(opname)
+
+                self.addRule(
+                    """
+                    dict_comp_async      ::= BUILD_MAP_0 LOAD_ARG
+                                             dict_comp_async
+
+                    expr                 ::= dict_comp_async
+                    expr                 ::= generator_exp_async
+                    expr                 ::= list_comp_async
+                    expr                 ::= set_comp_async
+
+                    func_async_middle   ::= POP_BLOCK JUMP_FORWARD COME_FROM_EXCEPT
+                                            DUP_TOP LOAD_GLOBAL COMPARE_OP POP_JUMP_IF_TRUE
+                                            END_FINALLY _come_froms
+
+                    # async_iter         ::= block_break SETUP_EXCEPT GET_ANEXT LOAD_CONST YIELD_FROM
+
+                    get_aiter            ::= expr GET_AITER
+
+                    list_afor            ::= get_aiter list_afor2
+
+                    list_comp_async      ::= BUILD_LIST_0 LOAD_ARG list_afor2
+                    list_iter            ::= list_afor
+
+
+                    set_afor             ::= get_aiter set_afor2
+                    set_iter             ::= set_afor
+                    set_iter             ::= set_for
+
+                    set_comp_async       ::= BUILD_SET_0 LOAD_ARG
+                                             set_comp_async
+
+                   """,
+                    nop_func,
+                )
+                custom_ops_processed.add(opname)
+
+
             elif opname == "GET_ANEXT":
                 self.addRule(
                     """

uncompyle6/parsers/parse37base.py

@@ -660,7 +660,7 @@ class Python37BaseParser(PythonParser):
                                             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
+                    list_comp_async     ::= BUILD_LIST_0 LOAD_ARG list_afor2
                     get_aiter           ::= expr GET_AITER
                     list_afor           ::= get_aiter list_afor2
                     list_iter           ::= list_afor

uncompyle6/scanners/scanner3.py

@@ -601,6 +601,10 @@ class Scanner3(Scanner):
                     # other parts like n_LOAD_CONST in pysource.py for example.
                     pattr = const
                     pass
+            elif opname == "LOAD_FAST" and argval == ".0":
+                # Used as the parameter of a list expression
+                opname = "LOAD_ARG"
+
             elif opname in ("MAKE_FUNCTION", "MAKE_CLOSURE"):
                 if self.version >= (3, 6):
                     # 3.6+ doesn't have MAKE_CLOSURE, so opname == 'MAKE_FUNCTION'

uncompyle6/semantics/customize36.py

@@ -693,6 +693,16 @@ def customize_for_version36(self, version):
 
     self.n_joined_str = n_joined_str
 
+    def n_list_comp_async(node):
+        self.write("[")
+        if node[0].kind == "load_closure":
+            self.listcomp_closure3(node)
+        else:
+            self.comprehension_walk_newer(node, iter_index=3, code_index=0)
+        self.write("]")
+        self.prune()
+    self.n_list_comp_async = n_list_comp_async
+
     # def kwargs_only_36(node):
     #     keys = node[-1].attr
     #     num_kwargs = len(keys)