Fix support of augmented assignment with no rvalue. Fixes GH-1.

progval · progval · commit 0960623312f9 · 2018-10-28T15:14:55.000+01:00
diff --git a/src/ast.rs b/src/ast.rs
@@ -327,6 +327,8 @@ pub enum Statement {
     Expressions(Vec<Expression>),
     // `lhs = rhs1 = rhs2` -> `lhs, vec![rhs1, rhs2]`
     Assignment(Vec<Expression>, Vec<Vec<Expression>>),
+    // `lhs: type` -> `lhs, type`
+    TypeAnnotation(Vec<Expression>, Expression),
     // `lhs: type = rhs` -> `lhs, type, rhs`
     TypedAssignment(Vec<Expression>, Expression, Vec<Expression>),
     // `lhs += rhs` -> `lhs, AugAssignOp::Add, rhs`
diff --git a/src/statements.rs b/src/statements.rs
@@ -64,13 +64,17 @@ named!(expr_stmt<StrSpan, Statement>,
   do_parse!(
     lhs: call!(ExpressionParser::<NewlinesAreNotSpaces>::testlist_star_expr) >>
     r: ws_nonl!(alt!(
-      // Case 1: "foo: bar = baz"
+      // Case 1: "foo: bar = baz" and "foo: bar"
       do_parse!(
         char!(':') >>
         typed: call!(ExpressionParser::<NewlinesAreNotSpaces>::test) >>
-        char!('=') >>
-        rhs: call!(ExpressionParser::<NewlinesAreNotSpaces>::test) >> (
-          Statement::TypedAssignment(lhs.clone(), *typed, vec![*rhs])
+        rhs: opt!(ws_nonl!(preceded!(char!('='),
+            call!(ExpressionParser::<NewlinesAreNotSpaces>::test)
+        ))) >> (
+          match rhs {
+              None => Statement::TypeAnnotation(lhs.clone(), *typed),
+              Some(rhs) => Statement::TypedAssignment(lhs.clone(), *typed, vec![*rhs]),
+          }
         )
       )
 
@@ -974,6 +978,18 @@ mod tests {
         )));
     }
 
+    #[test]
+    fn test_typeannotation() {
+        assert_parse_eq(small_stmt(make_strspan("foo: bar")), Ok((make_strspan(""),
+            Statement::TypeAnnotation(
+                vec![
+                    Expression::Name("foo".to_string()),
+                ],
+                Expression::Name("bar".to_string()),
+            )
+        )));
+    }
+
     #[test]
     fn test_augassign() {
         assert_parse_eq(small_stmt(make_strspan("foo:bar = baz")), Ok((make_strspan(""),
diff --git a/src/visitors/printer.rs b/src/visitors/printer.rs
@@ -113,6 +113,11 @@ fn format_statement(indent: usize, stmt: &Statement) -> String {
             }
             s.push_str("\n");
         },
+        Statement::TypeAnnotation(ref lhs, ref typed) => {
+            s.push_str(&format!("{}: {}\n",
+                comma_join(lhs.iter().map(format_expr)),
+                format_expr(typed)));
+        },
         Statement::TypedAssignment(ref lhs, ref typed, ref rhs) => {
             s.push_str(&format!("{}:{} = {}\n",
                 comma_join(lhs.iter().map(format_expr)),
