Add generator expression syntax

windelbouwman · windelbouwman · commit bc23e60962da · 2018-10-20T12:16:18.000+02:00
diff --git a/parser/src/ast.rs b/parser/src/ast.rs
@@ -210,6 +210,7 @@ pub struct Parameters {
 
 #[derive(Debug, PartialEq)]
 pub enum ComprehensionKind {
+    GeneratorExpression { element: Expression },
     List { element: Expression },
     Set { element: Expression },
     Dict { key: Expression, value: Expression },
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
@@ -707,6 +707,12 @@ Atom: ast::Expression = {
          }
      }
   },
+  "(" <e:Test> <c:CompFor> ")" => {
+        ast::Expression::Comprehension {
+            kind: Box::new(ast::ComprehensionKind::GeneratorExpression { element: e }),
+            generators: c,
+        }
+  },
   "{" <e:TestDict?> "}" => ast::Expression::Dict { elements: e.unwrap_or(Vec::new()) },
   "{" <e:TestDictComp> "}" => e,
   "{" <e:TestSet> "}" => ast::Expression::Set { elements: e },
@@ -862,7 +868,16 @@ ArgumentList: (Vec<ast::Expression>, Vec<ast::Keyword>) = {
 };
 
 FunctionArgument: (Option<String>, ast::Expression) = {
-    <e:Test> => (None, e),
+    <e:Test> <c:CompFor?> => {
+        let expr = match c {
+            Some(c) => ast::Expression::Comprehension {
+                kind: Box::new(ast::ComprehensionKind::GeneratorExpression { element: e }),
+                generators: c,
+            },
+            None => e,
+        };
+        (None, expr)
+    },
     <i:Identifier> "=" <e:Test> => (Some(i.clone()), e),
 };
 
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
@@ -933,6 +933,7 @@ impl Compiler {
         assert!(generators.len() > 0);
 
         let name = match kind {
+            ast::ComprehensionKind::GeneratorExpression { .. } => "<genexpr>",
             ast::ComprehensionKind::List { .. } => "<listcomp>",
             ast::ComprehensionKind::Set { .. } => "<setcomp>",
             ast::ComprehensionKind::Dict { .. } => "<dictcomp>",
@@ -948,6 +949,7 @@ impl Compiler {
 
         // Create empty object of proper type:
         match kind {
+            ast::ComprehensionKind::GeneratorExpression { .. } => {}
             ast::ComprehensionKind::List { .. } => {
                 self.emit(Instruction::BuildList { size: 0 });
             }
@@ -990,30 +992,26 @@ impl Compiler {
         }
 
         match kind {
+            ast::ComprehensionKind::GeneratorExpression { element } => {
+                self.compile_expression(element)?;
+                self.emit(Instruction::YieldValue);
+            }
             ast::ComprehensionKind::List { element } => {
-                // Evaluate element:
                 self.compile_expression(element)?;
-
-                // List append:
                 self.emit(Instruction::ListAppend {
                     i: 1 + generators.len(),
                 });
             }
             ast::ComprehensionKind::Set { element } => {
-                // Evaluate element:
                 self.compile_expression(element)?;
-
-                // List append:
                 self.emit(Instruction::SetAdd {
                     i: 1 + generators.len(),
                 });
             }
             ast::ComprehensionKind::Dict { key, value } => {
-                // Evaluate value and element:
                 self.compile_expression(value)?;
                 self.compile_expression(key)?;
 
-                // List append:
                 self.emit(Instruction::MapAdd {
                     i: 1 + generators.len(),
                 });
