Fix python 3 set comprehension and ...

Add a few set/list comprehension offsets for Python 3

pytest/test_deparse.py

@@ -20,6 +20,14 @@ def for_range_stmt():
     for i in range(2):
         i+1
 
+# FIXME: add this test
+def set_comp():
+    {y for y in range(3)}
+
+# FIXME: add this test
+def list_comp():
+    [y for y in range(3)]
+
 def get_parsed_for_fn(fn):
     code = fn.__code__ if PYTHON3 else fn.func_code
     return deparse(PYTHON_VERSION, code)

test/simple_source/comprehension/05_set_comprehension.py

@@ -0,0 +1,2 @@
+# Bug in python 3.x handling set comprehensions
+{y for y in range(3)}

uncompyle6/parser.py

@@ -50,7 +50,9 @@ class PythonParser(GenericASTBuilder):
 
         if nt in collect and len(args) > 1:
             #
-            #  Collect iterated thingies together.
+            #  Collect iterated thingies together. That is rather than
+            #  stmts -> stmts stmt -> stmts stmt -> ...
+            #  stmms -> stmt stmt ...
             #
             rv = args[0]
             rv.append(args[1])

uncompyle6/parsers/parse3.py

@@ -494,15 +494,14 @@ class Python3Parser(PythonParser):
                     rule = ("listcomp ::= LOAD_LISTCOMP MAKE_FUNCTION_0 expr "
                             "GET_ITER CALL_FUNCTION_1")
                 self.add_unique_rule(rule, opname, token.attr, customize)
-            # FIXME: add in after fixing bug in semantics
+            elif opname == 'LOAD_SETCOMP':
-            # elif opname == 'LOAD_SETCOMP':
+                if self.version >= 3.4:
-            #     if self.version >= 3.4:
+                    rule = ("setcomp ::= LOAD_SETCOMP LOAD_CONST MAKE_FUNCTION_0 expr "
-            #         rule = ("setcomp ::= LOAD_SETCOMP LOAD_CONST MAKE_FUNCTION_0 expr "
+                            "GET_ITER CALL_FUNCTION_1")
-            #                 "GET_ITER CALL_FUNCTION_1")
+                else:
-            #     else:
+                    rule = ("setcomp ::= LOAD_SETCOMP MAKE_FUNCTION_0 expr "
-            #         rule = ("setcomp ::= LOAD_SETCOMP MAKE_FUNCTION_0 expr "
+                            "GET_ITER CALL_FUNCTION_1")
-            #                 "GET_ITER CALL_FUNCTION_1")
+                self.add_unique_rule(rule, opname, token.attr, customize)
-            #     self.add_unique_rule(rule, opname, token.attr, customize)
             elif opname == 'LOAD_BUILD_CLASS':
                 self.custom_build_class_rule(opname, i, token, tokens, customize)
             elif opname_base in ('BUILD_LIST', 'BUILD_TUPLE', 'BUILD_SET'):

uncompyle6/semantics/fragments.py

@@ -580,27 +580,43 @@ class FragmentsWalker(pysource.SourceWalker, object):
 
         ast = self.build_ast(code._tokens, code._customize)
         self.customize(code._customize)
-        ast = ast[0][0][0][0][0]
+        # skip over stmts sstmt smt
+        ast = ast[0][0][0]
 
-        n = ast[iter_index]
+        if ast == 'setcomp_func':
-        assert n == 'list_iter'
+            for k in ast:
+                if k == 'comp_iter':
+                    n = k
+                elif k == 'designator':
+                    designator = k
+                    pass
+                pass
+            pass
+        else:
+            ast = ast[0][0]
+            n = ast[iter_index]
+            assert n == 'list_iter'
+
+        ## FIXME: I'm not totally sure this is right.
 
         # find innermost node
-        while n == 'list_iter':
+        list_if_node = None
+        while n in ('list_iter', 'comp_iter'):
             n = n[0] # recurse one step
             if   n == 'list_for':
-                designator = n[2]
+                if n[2] == 'designator':
+                    designator = n[2]
                 n = n[3]
-            elif n == 'list_if':
+            elif n in ['list_if', 'list_if_not']:
-                # FIXME: just a guess
+                list_if_node = n[0]
-                designator = n[1]
+                if n[1] == 'designator':
+                    designator = n[1]
+                n = n[2]
+                pass
+            pass
 
-                n = n[2]
+        assert n.type in ('lc_body', 'comp_body'), ast
-            elif n == 'list_if_not':
+        assert designator, "Couldn't find designator in list/set comprehension"
-                # FIXME: just a guess
-                designator = n[1]
-                n = n[2]
-        assert n == 'lc_body', ast
 
         self.preorder(n[0])
         self.write(' for ')
@@ -612,7 +628,9 @@ class FragmentsWalker(pysource.SourceWalker, object):
         node[-3].parent = node
         self.preorder(node[-3])
         self.set_pos_info(node[-3], start, len(self.f.getvalue()))
-        # self.preorder(ast[iter_index])
+        if list_if_node:
+            self.write(' if ')
+            self.preorder(list_if_node)
         self.prec = p
 
     def listcomprehension_walk2(self, node):
@@ -680,7 +698,12 @@ class FragmentsWalker(pysource.SourceWalker, object):
     def n_setcomp(self, node):
         start = len(self.f.getvalue())
         self.write('{')
-        self.comprehension_walk(node, 4)
+        if node[0] == 'LOAD_SETCOMP':
+            start = len(self.f.getvalue())
+            self.set_pos_info(node[0], start-1, start)
+            self.listcomprehension_walk3(node, 1, 0)
+        else:
+            self.comprehension_walk(node, iter_index=4)
         self.write('}')
         self.set_pos_info(node, start, len(self.f.getvalue()))
         self.prune()
@@ -767,7 +790,7 @@ class FragmentsWalker(pysource.SourceWalker, object):
         self.prune()
 
     def gen_source(self, ast, name, customize, isLambda=False, returnNone=False):
-        """convert AST to source code"""
+        """convert AST to Python source code"""
 
         rn = self.return_none
         self.return_none = returnNone
@@ -1585,7 +1608,7 @@ if __name__ == '__main__':
         return fn.__code__
 
     def test():
-        [x for x in range(3)]
+        {x for x in range(3)}
 
     def gcd(a, b):
         if a > b:

uncompyle6/semantics/pysource.py

@@ -590,7 +590,6 @@ class SourceWalker(GenericASTTraversal, object):
     def write(self, *data):
         if (len(data) == 0) or (len(data) == 1 and data[0] == ''):
             return
-#        import pdb; pdb.set_trace()
         out = ''.join((str(j) for j in data))
         n = 0
         for i in out:
@@ -812,10 +811,6 @@ class SourceWalker(GenericASTTraversal, object):
             if node[-2][0][-1] != 'BUILD_TUPLE_0':
                 node[-2][0].type = 'build_tuple2'
         self.default(node)
-#        maybe_tuple = node[-2][-1]
-#        if maybe_tuple.type.startswith('BUILD_TUPLE'):
-#            maybe_tuple.type = 'build_tuple2'
-#        self.default(node)
 
     n_store_subscr = n_binary_subscr = n_delete_subscr
 
@@ -988,7 +983,7 @@ class SourceWalker(GenericASTTraversal, object):
         self.prune() # stop recursing
 
     def n_list_compr(self, node):
-        """List comprehensions the way they are done in Python2.
+        """List comprehensions the way they are done in Python 2.
         """
         p = self.prec
         self.prec = 27
@@ -1057,7 +1052,10 @@ class SourceWalker(GenericASTTraversal, object):
 
     def n_setcomp(self, node):
         self.write('{')
-        self.comprehension_walk(node, iter_index=4)
+        if node[0] == 'LOAD_SETCOMP':
+            self.listcomprehension_walk3(node, 1, 0)
+        else:
+            self.comprehension_walk(node, iter_index=4)
         self.write('}')
         self.prune()
 
@@ -1076,17 +1074,28 @@ class SourceWalker(GenericASTTraversal, object):
 
         ast = self.build_ast(code._tokens, code._customize)
         self.customize(code._customize)
-        ast = ast[0][0][0][0][0]
+        # skip over stmts sstmt smt
-
+        ast = ast[0][0][0]
-        n = ast[iter_index]
+        designator = None
-        assert n == 'list_iter'
+        if ast == 'setcomp_func':
+            for k in ast:
+                if k == 'comp_iter':
+                    n = k
+                elif k == 'designator':
+                    designator = k
+                    pass
+                pass
+            pass
+        else:
+            ast = ast[0][0]
+            n = ast[iter_index]
+            assert n == 'list_iter'
 
         ## FIXME: I'm not totally sure this is right.
 
         # find innermost node
-        designator = None
         list_if_node = None
-        while n == 'list_iter':
+        while n in ('list_iter', 'comp_iter'):
             n = n[0] # recurse one step
             if   n == 'list_for':
                 if n[2] == 'designator':
@@ -1099,8 +1108,8 @@ class SourceWalker(GenericASTTraversal, object):
                 n = n[2]
                 pass
             pass
-        assert n == 'lc_body', ast
+        assert n.type in ('lc_body', 'comp_body'), ast
-        assert designator, "Couldn't find designator in list comprehension"
+        assert designator, "Couldn't find designator in list/set comprehension"
 
         self.preorder(n[0])
         self.write(' for ')
@@ -1322,10 +1331,7 @@ class SourceWalker(GenericASTTraversal, object):
                     l = list(kv_node)
                 i = 0
                 while i < len(l):
-                    try:
+                    name = self.traverse(l[i+1], indent='')
-                        name = self.traverse(l[i+1], indent='')
-                    except:
-                        from trepan.api import debug; debug()
                     value = self.traverse(l[i], indent=self.indent+(len(name)+2)*' ')
                     self.write(sep, name, ': ', value)
                     sep = line_seperator
@@ -1504,7 +1510,6 @@ class SourceWalker(GenericASTTraversal, object):
             elif typ == 'C':
                 low, high, sep = entry[arg]
                 remaining = len(node[low:high])
-                # remaining = len(node[low:high])
                 for subnode in node[low:high]:
                     self.preorder(subnode)
                     remaining -= 1
@@ -1821,7 +1826,7 @@ class SourceWalker(GenericASTTraversal, object):
         self.classes.pop(-1)
 
     def gen_source(self, ast, name, customize, isLambda=False, returnNone=False):
-        """convert AST to source code"""
+        """convert AST to Python source code"""
 
         rn = self.return_none
         self.return_none = returnNone