comprehension starred expression compatibility

youknowone · youknowone · commit e4be37633698 · 2020-05-30T01:07:38.000Z
diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
@@ -2027,21 +2027,33 @@ impl<O: OutputStream> Compiler<O> {
             }
         }
 
+        let mut compile_element = |element| {
+            self.compile_expression(element).map_err(|e| {
+                if matches!(e.error, CompileErrorType::InvalidStarExpr) {
+                    self.error(CompileErrorType::SyntaxError(
+                        "iterable unpacking cannot be used in comprehension".to_owned(),
+                    ))
+                } else {
+                    e
+                }
+            })
+        };
+
         match kind {
             ast::ComprehensionKind::GeneratorExpression { element } => {
-                self.compile_expression(element)?;
+                compile_element(element)?;
                 self.mark_generator();
                 self.emit(Instruction::YieldValue);
                 self.emit(Instruction::Pop);
             }
             ast::ComprehensionKind::List { element } => {
-                self.compile_expression(element)?;
+                compile_element(element)?;
                 self.emit(Instruction::ListAppend {
                     i: 1 + generators.len(),
                 });
             }
             ast::ComprehensionKind::Set { element } => {
-                self.compile_expression(element)?;
+                compile_element(element)?;
                 self.emit(Instruction::SetAdd {
                     i: 1 + generators.len(),
                 });
