Allow tuple unpacking in return and yield statements (bpo-32117)

progval · progval · commit 935227f5cc2c · 2018-10-28T15:07:24.000+01:00
fd97d1f1af910a6222ea12aec42c456b64f9aee4
diff --git a/src/expressions.rs b/src/expressions.rs
@@ -375,6 +375,28 @@ named!(pub possibly_empty_testlist<StrSpan, Vec<Expression>>,
   )
 );
 
+// testlist_star_expr: (test|star_expr) (',' (test|star_expr))* [',']
+named!(pub testlist_star_expr<StrSpan, Vec<Expression>>,
+  do_parse!(
+    list: separated_nonempty_list!(
+      ws_auto!(char!(',')),
+      map!(alt!(
+        call!(Self::test)
+      | call!(Self::star_expr)
+      ), |e| *e)
+    ) >>
+    trailing_comma: opt!(ws_auto!(char!(','))) >> (
+      if trailing_comma.is_some() && list.len() < 2 {
+          // This prevents "foo, =" from being parsed as "foo ="
+          vec![Expression::TupleLiteral(list.into_iter().map(SetItem::Unique).collect())]
+      }
+      else {
+          list
+      }
+    )
+  )
+);
+
 } // end ExpressionParser
 
 /*********************************************************************
@@ -568,13 +590,13 @@ named_args!(comp_if(acc: Vec<ComprehensionChunk>) <StrSpan, Vec<ComprehensionChu
 
 
 // yield_expr: 'yield' [yield_arg]
-// yield_arg: 'from' test | testlist
+// yield_arg: 'from' test | testlist_star_expr
 named!(pub yield_expr<StrSpan, Expression>,
   ws_auto!(preceded!(
     keyword!("yield"),
     ws_auto!(alt!(
       preceded!(ws_auto!(keyword!("from")), call!(Self::test)) => { |e| Expression::YieldFrom(e) }
-    | call!(Self::testlist) => { |e| Expression::Yield(e) }
+    | call!(Self::testlist_star_expr) => { |e| Expression::Yield(e) }
     | tag!("") => { |_| Expression::Yield(Vec::new()) }
     ))
   ))
@@ -1652,4 +1674,16 @@ mod tests {
             ))
         )));
     }
+
+    #[test]
+    fn test_unpack() {
+        let testlist_star_expr = ExpressionParser::<NewlinesAreNotSpaces>::testlist_star_expr;
+        assert_parse_eq(testlist_star_expr(make_strspan("foo,")), Ok((make_strspan(""),
+            vec![
+                Expression::TupleLiteral(vec![
+                    SetItem::Unique(Expression::Name("foo".to_string())),
+                ]),
+            ]
+        )));
+    }
 }
diff --git a/src/statements.rs b/src/statements.rs
@@ -62,7 +62,7 @@ named!(small_stmt<StrSpan, Statement>,
 // annassign: ':' test ['=' test]
 named!(expr_stmt<StrSpan, Statement>,
   do_parse!(
-    lhs: testlist_star_expr >>
+    lhs: call!(ExpressionParser::<NewlinesAreNotSpaces>::testlist_star_expr) >>
     r: ws_nonl!(alt!(
       // Case 1: "foo: bar = baz"
       do_parse!(
@@ -89,7 +89,7 @@ named!(expr_stmt<StrSpan, Statement>,
       do_parse!(
         rhs: many0!(ws_nonl!(preceded!(char!('='), alt!(
           call!(ExpressionParser::<NewlinesAreNotSpaces>::yield_expr) => { |e| vec![e] }
-        | testlist_star_expr
+        | call!(ExpressionParser::<NewlinesAreNotSpaces>::testlist_star_expr)
         )))) >> (
           Statement::Assignment(lhs, rhs)
         )
@@ -99,28 +99,6 @@ named!(expr_stmt<StrSpan, Statement>,
   )
 );
 
-// testlist_star_expr: (test|star_expr) (',' (test|star_expr))* [',']
-named!(testlist_star_expr<StrSpan, Vec<Expression>>,
-  do_parse!(
-    list: separated_nonempty_list!(
-      ws_nonl!(char!(',')),
-      map!(alt!(
-        call!(ExpressionParser::<NewlinesAreNotSpaces>::test)
-      | call!(ExpressionParser::<NewlinesAreNotSpaces>::star_expr)
-      ), |e| *e)
-    ) >>
-    trailing_comma: opt!(ws_nonl!(char!(','))) >> (
-      if trailing_comma.is_some() && list.len() < 2 {
-          // This prevents "foo, =" from being parsed as "foo ="
-          vec![Expression::TupleLiteral(list.into_iter().map(SetItem::Unique).collect())]
-      }
-      else {
-          list
-      }
-    )
-  )
-);
-
 // augassign: ('+=' | '-=' | '*=' | '@=' | '/=' | '%=' | '&=' | '|=' | '^=' |
 //            '<<=' | '>>=' | '**=' | '//=')
 named!(augassign<StrSpan, AugAssignOp>,
@@ -159,15 +137,15 @@ named!(pass_stmt<StrSpan, Statement>,
 // flow_stmt: break_stmt | continue_stmt | return_stmt | raise_stmt | yield_stmt
 // break_stmt: 'break'
 // continue_stmt: 'continue'
-// return_stmt: 'return' [testlist]
+// return_stmt: 'return' [testlist_star_expr]
 // yield_stmt: yield_expr
 named!(flow_stmt<StrSpan, Statement>,
   alt!(
     keyword!("break") => { |_| Statement::Break }
   | keyword!("continue") => { |_| Statement::Continue }
   | preceded!(
       tuple!(keyword!("return"), spaces_nonl),
-      return_error!(ws_nonl!(call!(ExpressionParser::<NewlinesAreNotSpaces>::possibly_empty_testlist)))
+      return_error!(ws_nonl!(call!(ExpressionParser::<NewlinesAreNotSpaces>::testlist_star_expr)))
     ) => { |e| Statement::Return(e) }
   | raise_stmt
   | call!(ExpressionParser::<NewlinesAreNotSpaces>::yield_expr)
@@ -1192,15 +1170,4 @@ mod tests {
             })]
         )));
     }
-
-    #[test]
-    fn test_unpack() {
-        assert_parse_eq(testlist_star_expr(make_strspan("foo,")), Ok((make_strspan(""),
-            vec![
-                Expression::TupleLiteral(vec![
-                    SetItem::Unique(Expression::Name("foo".to_string())),
-                ]),
-            ]
-        )));
-    }
 }
