Port Python 2.7 list comprehension fixes by @reductor from #78

zrax · Oct 4, 2019 · b5696c8 · b5696c8
1 parent 314f4a1
commit b5696c8
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 2 deletions.
diff --git a/ASTNode.h b/ASTNode.h
@@ -497,14 +497,17 @@ class ASTIterBlock : public ASTBlock {
 
     PycRef<ASTNode> iter() const { return m_iter; }
     PycRef<ASTNode> index() const { return m_idx; }
+    PycRef<ASTNode> condition() const { return m_cond; }
     bool isComprehension() const { return m_comp; }
 
     void setIndex(PycRef<ASTNode> idx) { m_idx = idx; init(); }
+    void setCondition(PycRef<ASTNode> cond) { m_cond = cond; }
     void setComprehension(bool comp) { m_comp = comp; }
 
 private:
     PycRef<ASTNode> m_iter;
     PycRef<ASTNode> m_idx;
+    PycRef<ASTNode> m_cond;
     bool m_comp;
 };
 

diff --git a/ASTree.cpp b/ASTree.cpp
@@ -272,7 +272,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
             }
             break;
         case Pyc::BUILD_MAP_A:
-            if (mod->majorVer() > 3 || (mod->majorVer() == 3 && mod->minorVer() >= 5)) {
+            if (mod->verCompare(3, 5) >= 0) {
                 auto map = new ASTMap;
                 for (int i=0; i<operand; ++i) {
                     PycRef<ASTNode> value = stack.top();
@@ -1029,6 +1029,15 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                         newcond = new ASTBinary(cond1, cond, ASTBinary::BIN_LOG_OR);
                     }
                     ifblk = new ASTCondBlock(top->blktype(), offs, newcond, neg);
+                } else if (curblock->blktype() == ASTBlock::BLK_FOR
+                            && curblock.cast<ASTIterBlock>()->isComprehension()
+                            && mod->verCompare(2, 7) >= 0) {
+                    /* Comprehension condition */
+                    curblock.cast<ASTIterBlock>()->setCondition(cond);
+                    stack_hist.pop();
+                    // TODO: Handle older python versions, where condition
+                    // is laid out a little differently.
+                    break;
                 } else {
                     /* Plain old if statement */
                     ifblk = new ASTCondBlock(ASTBlock::BLK_IF, offs, cond, neg);
@@ -1056,7 +1065,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
 
                         blocks.pop();
                         curblock = blocks.top();
-                    } if (curblock->blktype() == ASTBlock::BLK_ELSE) {
+                    } else if (curblock->blktype() == ASTBlock::BLK_ELSE) {
                         stack = stack_hist.top();
                         stack_hist.pop();
 
@@ -1271,6 +1280,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
 
                 if (curblock->blktype() == ASTBlock::BLK_FOR
                         && curblock.cast<ASTIterBlock>()->isComprehension()) {
+                    stack.pop();
                     stack.push(new ASTComprehension(value));
                 } else {
                     stack.push(new ASTSubscr(list, value)); /* Total hack */
@@ -2333,6 +2343,10 @@ void print_src(PycRef<ASTNode> node, PycModule* mod)
                 print_src(gen->index(), mod);
                 fputs(" in ", pyc_output);
                 print_src(gen->iter(), mod);
+                if (gen->condition()) {
+                    fprintf(pyc_output, " if ");
+                    print_src(gen->condition(), mod);
+                }
             }
             fputs(" ]", pyc_output);
         }

diff --git a/tests/compiled/test_listComprehensions.2.7.pyc b/tests/compiled/test_listComprehensions.2.7.pyc
diff --git a/tests/tokenized/test_listComprehensions.txt b/tests/tokenized/test_listComprehensions.txt
@@ -0,0 +1,21 @@
+XXX = range ( 4 ) <EOL>
+print [ i for i in XXX ] <EOL>
+print <EOL>
+print [ i for i in ( 1 , 2 , 3 , 4 ) ] <EOL>
+print <EOL>
+print [ ( i , 1 ) for i in XXX ] <EOL>
+print <EOL>
+print [ i * 2 for i in range ( 4 ) ] <EOL>
+print <EOL>
+print [ i * j for i in range ( 4 ) for j in range ( 7 ) ] <EOL>
+print [ i * 2 for i in range ( 4 ) if i == 0 ] <EOL>
+print [ ( i , i ** 2 ) for i in range ( 4 ) if i % 2 == 0 ] <EOL>
+print [ i * j for i in range ( 4 ) if i == 2 for j in range ( 7 ) if i + i % 2 == 0 ] <EOL>
+seq1 = 'abc' <EOL>
+seq2 = ( 1 , 2 , 3 ) <EOL>
+[ ( x , y ) for x in seq1 for y in seq2 ] <EOL>
+def flatten ( seq ) : <EOL>
+<INDENT>
+return [ x for subseq in seq for x in subseq ] <EOL>
+<OUTDENT>
+print flatten ( [ [ 0 ] , [ 1 , 2 , 3 ] , [ 4 , 5 ] , [ 6 , 7 , 8 , 9 ] , [ ] ] ) <EOL>