From 698ae764f1ffd6a6bfaf3ddc275cd93809fe196c Mon Sep 17 00:00:00 2001
From: Batuhan Taskaya <batuhanosmantaskaya@gmail.com>
Date: Mon, 30 Nov 2020 03:11:09 +0300
Subject: [PATCH] Handle raw yield expressions as RHS for augmented/annotated
 assignments

---
 parser/src/python.lalrpop | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
index f3a4fb642..e9b9aaaea 100644
--- a/parser/src/python.lalrpop
+++ b/parser/src/python.lalrpop
@@ -107,7 +107,7 @@ ExpressionStatement: ast::Statement = {
             }
         }
     },
-    <location:@L> <target:TestOrStarExprList> <op:AugAssign> <rhs:TestList> => {
+    <location:@L> <target:TestOrStarExprList> <op:AugAssign> <rhs:TestListOrYieldExpr> => {
         ast::Statement {
             location,
             node: ast::StatementType::AugAssign {
@@ -117,23 +117,27 @@ ExpressionStatement: ast::Statement = {
             },
         }
     },
-    <location:@L> <target:Test> ":" <annotation:Test> <rhs:("=" Test)?> => {
+    <location:@L> <target:Test> ":" <annotation:Test> <rhs:AssignSuffix?> => {
         ast::Statement {
             location,
             node: ast::StatementType::AnnAssign {
                 target: Box::new(target),
                 annotation: Box::new(annotation),
-                value: rhs.map(|e| e.1)
+                value: rhs
             },
         }
     },
 };
 
 AssignSuffix: ast::Expression = {
-    "=" <e:TestList> => e,
-    "=" <e:YieldExpr> => e,
+    "=" <e:TestListOrYieldExpr> => e
 };
 
+TestListOrYieldExpr: ast::Expression = {
+    TestList,
+    YieldExpr
+}
+
 TestOrStarExprList: ast::Expression = {
     <location:@L> <elements:OneOrMore<TestOrStarExpr>> <comma:","?> => {
         if elements.len() == 1 && comma.is_none() {