From 29ce0cf2e7e2edaebfcd0c9b0317d378d6686b43 Mon Sep 17 00:00:00 2001
From: Windel Bouwman <windel@windel.nl>
Date: Sat, 27 Oct 2018 11:35:46 +0200
Subject: [PATCH] Add support for if-filters in list comprehensions.

---
 parser/src/parser.rs             | 23 +++++++++++++--
 parser/src/python.lalrpop        | 48 ++++++++------------------------
 tests/snippets/comprehensions.py |  7 ++---
 vm/src/compile.rs                | 10 +++++++
 4 files changed, 45 insertions(+), 43 deletions(-)

diff --git a/parser/src/parser.rs b/parser/src/parser.rs
index cd493e6fc..f1b1afd8f 100644
--- a/parser/src/parser.rs
+++ b/parser/src/parser.rs
@@ -381,7 +381,7 @@ mod tests {
 
     #[test]
     fn test_parse_double_list_comprehension() {
-        let source = String::from("[x for y, y2 in z for a in b]");
+        let source = String::from("[x for y, y2 in z for a in b if a < 5 if a > 10]");
         let parse_ast = parse_expression(&source).unwrap();
         assert_eq!(
             parse_ast,
@@ -415,7 +415,26 @@ mod tests {
                         iter: ast::Expression::Identifier {
                             name: "b".to_string()
                         },
-                        ifs: vec![],
+                        ifs: vec![
+                            ast::Expression::Compare {
+                                a: Box::new(ast::Expression::Identifier {
+                                    name: "a".to_string()
+                                }),
+                                op: ast::Comparison::Less,
+                                b: Box::new(ast::Expression::Number {
+                                    value: ast::Number::Integer { value: 5 }
+                                }),
+                            },
+                            ast::Expression::Compare {
+                                a: Box::new(ast::Expression::Identifier {
+                                    name: "a".to_string()
+                                }),
+                                op: ast::Comparison::Greater,
+                                b: Box::new(ast::Expression::Number {
+                                    value: ast::Number::Integer { value: 10 }
+                                }),
+                            },
+                        ],
                     }
                 ],
             }
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
index d575f00ab..2641e9eb5 100644
--- a/parser/src/python.lalrpop
+++ b/parser/src/python.lalrpop
@@ -904,34 +904,16 @@ StarExpr: ast::Expression = {
 };
 
 // Comprehensions:
-CompIter: (Option<Vec<ast::Comprehension>>, Option<ast::Expression>) = {
-   //  <e:CompIf> => (None, Some(e)),
-    <c:CompFor> => (Some(c), None),
+CompFor: Vec<ast::Comprehension> = {
+    <c:SingleForComprehension+> => c,
 };
 
-CompFor: Vec<ast::Comprehension> = {
-    "for" <e:ExpressionList> "in" <i:OrTest> <c2:CompIter?> => {
-        match c2 {
-            None => {
-                vec![ast::Comprehension {
-                    target: e,
-                    iter: i,
-                    ifs: vec![],
-                }]
-            },
-            Some((Some(mut cphs), None)) => {
-                let mut res = vec![];
-                res.push(ast::Comprehension {
-                    target: e,
-                    iter: i,
-                    ifs: vec![],
-                });
-                res.append(&mut cphs);
-                res
-            },
-            _ => {
-                panic!("TODO");
-            },
+SingleForComprehension: ast::Comprehension = {
+    "for" <e:ExpressionList> "in" <i:OrTest> <c2:ComprehensionIf*> => {
+        ast::Comprehension {
+            target: e,
+            iter: i,
+            ifs: c2,
         }
     }
 };
@@ -940,17 +922,9 @@ ExpressionNoCond: ast::Expression = {
     OrTest,
 };
 
-//CompIf: ast::Expression = {
-//    "if" <c:ExpressionNoCond> <c2:CompIter?> => {
-//        match c2 {
-//            None => {
-//                vec![]
-//            },
-//            Some() => {
-//            },
-//        }
-//    }
-//};
+ComprehensionIf: ast::Expression = {
+    "if" <c:ExpressionNoCond> => c,
+};
 
 ArgumentList: (Vec<ast::Expression>, Vec<ast::Keyword>) = {
     <e: Comma<FunctionArgument>> => {
diff --git a/tests/snippets/comprehensions.py b/tests/snippets/comprehensions.py
index 4a4563dce..7fc137a5d 100644
--- a/tests/snippets/comprehensions.py
+++ b/tests/snippets/comprehensions.py
@@ -18,7 +18,6 @@ v = {b * 2 for b in x}
 # TODO:
 #u = {str(b): b-2 for b in x}
 
-# TODO: With if filtering:
-#y = [a+2 for a in x if a % 2]
-#print(y)
-#assert y == [3, 5]
+y = [a+2 for a in x if a % 2]
+print(y)
+assert y == [3, 5]
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
index cb6e1b568..7ddd2ef3e 100644
--- a/vm/src/compile.rs
+++ b/vm/src/compile.rs
@@ -1050,6 +1050,16 @@ impl Compiler {
             self.emit(Instruction::ForIter);
 
             self.compile_store(&generator.target)?;
+
+            // Now evaluate the ifs:
+            for if_condition in &generator.ifs {
+                self.compile_test(
+                    if_condition,
+                    None,
+                    Some(start_label),
+                    EvalContext::Statement,
+                )?
+            }
         }
 
         match kind {