diff --git a/parser/src/ast.rs b/parser/src/ast.rs index b134e44f7..21a38ea23 100644 --- a/parser/src/ast.rs +++ b/parser/src/ast.rs @@ -78,7 +78,7 @@ pub enum Statement { body: Vec, }, For { - target: Vec, + target: Expression, iter: Vec, body: Vec, orelse: Option>, @@ -200,7 +200,7 @@ pub enum Expression { #[derive(Debug, PartialEq, Clone)] pub struct Comprehension { - pub target: Vec, // TODO: should this be a non-vector? + pub target: Expression, pub iter: Expression, pub ifs: Vec, } diff --git a/parser/src/parser.rs b/parser/src/parser.rs index 73abab09f..58ab1f96d 100644 --- a/parser/src/parser.rs +++ b/parser/src/parser.rs @@ -337,9 +337,9 @@ mod tests { name: "x".to_string() }), generators: vec![ast::Comprehension { - target: vec![ast::Expression::Identifier { + target: ast::Expression::Identifier { name: "y".to_string() - },], + }, iter: ast::Expression::Identifier { name: "z".to_string() }, @@ -348,4 +348,45 @@ mod tests { } ); } + + #[test] + fn test_parse_double_list_comprehension() { + let source = String::from("[x for y, y2 in z for a in b]"); + let parse_ast = parse_expression(&source).unwrap(); + assert_eq!( + parse_ast, + ast::Expression::ListComprehension { + element: Box::new(ast::Expression::Identifier { + name: "x".to_string() + }), + generators: vec![ + ast::Comprehension { + target: ast::Expression::Tuple { + elements: vec![ + ast::Expression::Identifier { + name: "y".to_string() + }, + ast::Expression::Identifier { + name: "y2".to_string() + }, + ], + }, + iter: ast::Expression::Identifier { + name: "z".to_string() + }, + ifs: vec![], + }, + ast::Comprehension { + target: ast::Expression::Identifier { + name: "a".to_string() + }, + iter: ast::Expression::Identifier { + name: "b".to_string() + }, + ifs: vec![], + } + ], + } + ); + } } diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop index c3de0bb49..81ec73016 100644 --- a/parser/src/python.lalrpop +++ b/parser/src/python.lalrpop @@ -298,8 +298,11 @@ ForStatement: ast::LocatedStatement = { None => None, }; ast::LocatedStatement { - location: loc, - node: ast::Statement::For { target: e, iter: t, body: s, orelse: or_else }, + location: loc, + node: ast::Statement::For { + target: e, + iter: t, body: s, orelse: or_else + }, } }, }; @@ -622,8 +625,14 @@ DictEntry: (ast::Expression, ast::Expression) = { ":" => (e1, e2), }; -ExpressionList: Vec = { - > => e, +ExpressionList: ast::Expression = { + > => { + if e.len() == 1 { + e.into_iter().next().unwrap() + } else { + ast::Expression::Tuple { elements: e } + } + }, }; #[inline] @@ -653,12 +662,12 @@ CompFor: Vec = { }, Some((Some(mut cphs), None)) => { let mut res = vec![]; - res.append(&mut cphs); res.push(ast::Comprehension { target: e, iter: i, ifs: vec![], }); + res.append(&mut cphs); res }, _ => { diff --git a/tests/snippets/comprehensions.py b/tests/snippets/comprehensions.py index 240480bb3..7b36480e9 100644 --- a/tests/snippets/comprehensions.py +++ b/tests/snippets/comprehensions.py @@ -4,3 +4,9 @@ x = [1, 2, 3] y = [a*a+1 for a in x] assert y == [2, 5, 10] +z = [(b, c) for b in x for c in y] +# print(z) +assert z == [ + (1, 2), (1, 5), (1, 10), + (2, 2), (2, 5), (2, 10), + (3, 2), (3, 5), (3, 10)] diff --git a/vm/src/compile.rs b/vm/src/compile.rs index 1f88b1426..4a33ac62d 100644 --- a/vm/src/compile.rs +++ b/vm/src/compile.rs @@ -250,17 +250,7 @@ impl Compiler { self.emit(Instruction::ForIter); // Start of loop iteration, set targets: - // TODO: can we use compile_store here? - for t in target { - match t { - ast::Expression::Identifier { name } => { - self.emit(Instruction::StoreName { - name: name.to_string(), - }); - } - _ => panic!("Not impl"), - } - } + self.compile_store(target); // Body of loop: self.compile_statements(body); @@ -921,17 +911,7 @@ impl Compiler { self.set_label(start_label); self.emit(Instruction::ForIter); - // TODO: can we use compile_store here? - for t in &generator.target { - match t { - ast::Expression::Identifier { name } => { - self.emit(Instruction::StoreName { - name: name.to_string(), - }); - } - _ => panic!("Not impl"), - } - } + self.compile_store(&generator.target); // Evaluate element: self.compile_expression(element);