Improve 3.0 list comprehensions

test/simple_source/bug31/06_listcomp.py

@@ -0,0 +1,17 @@
+# Python 3.0 comprehensions can produce different code from
+# all other Python versions. Thanks, Python!
+
+# Adapted from 3.0 ast.py
+# This code is RUNNABLE!
+def _format(node):
+    return [(a, int(b)) for a, b in node.items()]
+
+# Adapted from 3.0 cmd.py
+def monthrange(ary, dotext):
+    return [a[3:] for a in ary if a.startswith(dotext)]
+
+x = {'a': '1', 'b': '2'}
+assert [('a', 1), ('b', 2)] == _format(x)
+
+ary = ["Monday", "Twoday", "Monmonth"]
+assert ['day', 'month'] == monthrange(ary, "Mon")

uncompyle6/parsers/parse30.py

@@ -61,6 +61,7 @@ class Python30Parser(Python31Parser):
 
         # JUMP_IF_TRUE POP_TOP as a replacement
         comp_if       ::= expr jmp_false comp_iter
+        comp_if       ::= expr jmp_false comp_iter JUMP_BACK POP_TOP
         comp_iter     ::= expr expr SET_ADD
         comp_iter     ::= expr expr LIST_APPEND
 

uncompyle6/semantics/pysource.py

@@ -1178,6 +1178,26 @@ class SourceWalker(GenericASTTraversal, object):
             stores = [ast[3]]
             assert ast[4] == 'comp_iter'
             n = ast[4]
+            # Find the list comprehension body. It is the inner-most
+            # node that is not comp_.. .
+            while n == 'comp_iter':
+                if n[0] == 'comp_for':
+                    n = n[0]
+                    stores.append(n[2])
+                    n = n[3]
+                elif n[0] in ('comp_if', 'comp_if_not'):
+                    n = n[0]
+                    # FIXME: just a guess
+                    if n[0].kind == 'expr':
+                        list_ifs.append(n)
+                    else:
+                        list_ifs.append([1])
+                    n = n[2]
+                    pass
+                else:
+                    break
+                pass
+
             # Skip over n[0] which is something like: _[1]
             self.preorder(n[1])