diff --git a/uncompyle6/semantics/customize3.py b/uncompyle6/semantics/customize3.py index 8fccfc7d..bb8abe27 100644 --- a/uncompyle6/semantics/customize3.py +++ b/uncompyle6/semantics/customize3.py @@ -154,6 +154,7 @@ def customize_for_version3(self, version): # recurse one step n = n[0] + # FIXME: adjust for set comprehension if n == "list_for": stores.append(n[2]) n = n[3] @@ -168,13 +169,12 @@ def customize_for_version3(self, version): c = c[0] collections.append(c) pass - elif n in ("list_if", "list_if_not"): - # FIXME: just a guess + elif n in ("list_if", "list_if_not", "list_if_or_not"): if n[0].kind == "expr": list_ifs.append(n) else: list_ifs.append([1]) - n = n[2] + n = n[-2] if n[-1] == "come_from_opt" else n[-1] pass elif n == "list_if37": list_ifs.append(n) @@ -184,7 +184,7 @@ def customize_for_version3(self, version): collections.append(n[0][0]) n = n[1] stores.append(n[1][0]) - n = n[3] + n = n[2] if n[2].kind == "list_iter" else n[3] pass assert n == "lc_body", ast diff --git a/uncompyle6/semantics/pysource.py b/uncompyle6/semantics/pysource.py index 632e2070..55dd8963 100644 --- a/uncompyle6/semantics/pysource.py +++ b/uncompyle6/semantics/pysource.py @@ -1183,7 +1183,6 @@ class SourceWalker(GenericASTTraversal, object): is_lambda = self.is_lambda if node[0].kind == "load_genexpr": self.is_lambda = False - self.closure_walk(node, collection_index=4) self.closure_walk(node, collection_index=4) self.is_lambda = is_lambda else: @@ -1416,31 +1415,79 @@ class SourceWalker(GenericASTTraversal, object): self.write("]") self.prune() + def get_comprehension_function(self, node, code_index: int): + """ + Build the body of a comprehension function and then + find the comprehension node buried in the tree which may + be surrounded with start-like symbols or dominiators,. + """ + self.prec = 27 + code_node = node[code_index] + if code_node == "load_genexpr": + code_node = code_node[0] + + code_obj = code_node.attr + assert iscode(code_obj), code_node + + code = Code(code_obj, self.scanner, self.currentclass, self.debug_opts["asm"]) + + # FIXME: is there a way we can avoid this? + # The problem is that in filterint top-level list comprehensions we can + # encounter comprehensions of other kinds, and lambdas + if self.compile_mode in ("listcomp",): # add other comprehensions to this list + p_save = self.p + self.p = get_python_parser( + self.version, compile_mode="exec", is_pypy=self.is_pypy, + ) + tree = self.build_ast( + code._tokens, code._customize, code, is_lambda=self.is_lambda + ) + self.p = p_save + else: + tree = self.build_ast( + code._tokens, code._customize, code, is_lambda=self.is_lambda + ) + + self.customize(code._customize) + + # skip over: sstmt, stmt, return, return_expr + # and other singleton derivations + if tree == "lambda_start": + if tree[0] in ("dom_start", "dom_start_opt"): + tree = tree[1] + + while len(tree) == 1 or ( + tree in ("stmt", "sstmt", "return", "return_expr", "return_expr_lambda") + ): + self.prec = 100 + tree = tree[1] if tree[0] in ("dom_start", "dom_start_opt") else tree[0] + return tree + def closure_walk(self, node, collection_index): """Dictionary and comprehensions using closure the way they are done in Python3. """ p = self.prec self.prec = 27 - code = Code(node[1].attr, self.scanner, self.currentclass) - ast = self.build_ast(code._tokens, code._customize, code) - self.customize(code._customize) + code_index = 0 if node[0] == "load_genexpr" else 1 + tree = self.get_comprehension_function(node, code_index=code_index) # Remove single reductions as in ("stmts", "sstmt"): - while len(ast) == 1: - ast = ast[0] + while len(tree) == 1: + tree = tree[0] - store = ast[3] + store = tree[3] collection = node[collection_index] - iter_index = 3 if ast == "genexpr_func_async" else 4 - n = ast[iter_index] + iter_index = 3 if tree == "genexpr_func_async" else 4 + n = tree[iter_index] list_if = None assert n == "comp_iter" # Find inner-most node. while n == "comp_iter": n = n[0] # recurse one step + # FIXME: adjust for set comprehension if n == "list_for": store = n[2] @@ -1459,7 +1506,7 @@ class SourceWalker(GenericASTTraversal, object): pass pass - assert n == "comp_body", ast + assert n == "comp_body", tree self.preorder(n[0]) self.write(" for ")