Annotated assignment syntax (#1432)

* Annotated assignment syntax

* Update generator script

* Rename to AnnotatedAssignStatement

Author: slozier

Src/IronPython/Compiler/Ast/AnnotatedAssignStatement.cs

@@ -0,0 +1,46 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information.
+
+#nullable enable
+
+using IronPython.Runtime.Binding;
+
+using MSAst = System.Linq.Expressions;
+
+namespace IronPython.Compiler.Ast {
+
+    public class AnnotatedAssignStatement : Statement {
+        public AnnotatedAssignStatement(Expression target, Expression annotation, Expression? value, bool simple) {
+            Target = target;
+            Annotation = annotation;
+            Value = value;
+            Simple = simple;
+        }
+
+        public Expression Target { get; }
+
+        public Expression Annotation { get; }
+
+        public Expression? Value { get; }
+
+        public bool Simple { get; }
+
+        public override MSAst.Expression Reduce() {
+            if (Value is null) {
+                return Empty();
+            } else {
+                return Target.TransformSet(Span, Value, PythonOperationKind.None);
+            }
+        }
+
+        public override void Walk(PythonWalker walker) {
+            if (walker.Walk(this)) {
+                Target.Walk(walker);
+                Annotation.Walk(walker);
+                Value?.Walk(walker);
+            }
+            walker.PostWalk(this);
+        }
+    }
+}

Src/IronPython/Compiler/Ast/PythonNameBinder.cs

@@ -189,6 +189,14 @@ public override bool Walk(AssignmentStatement node) {
             return true;
         }
 
+        // AnnotatedAssignStatement
+        public override bool Walk(AnnotatedAssignStatement node) {
+            node.Parent = _currentScope;
+            node.Target.Walk(_define);
+            return true;
+        }
+
+        // AugmentedAssignStatement
         public override bool Walk(AugmentedAssignStatement node) {
             node.Parent = _currentScope;
             node.Left.Walk(_define);
@@ -727,7 +735,6 @@ public override bool Walk(NameExpression node) {
             return true;
         }
 
-
         // NonlocalStatement
         public override bool Walk(NonlocalStatement node) {
             node.Parent = _currentScope;

Src/IronPython/Compiler/Ast/PythonWalker.Generated.cs

@@ -124,6 +124,10 @@ public virtual void PostWalk(YieldExpression node) { }
         public virtual bool Walk(YieldFromExpression node) { return true; }
         public virtual void PostWalk(YieldFromExpression node) { }
 
+        // AnnotatedAssignStatement
+        public virtual bool Walk(AnnotatedAssignStatement node) { return true; }
+        public virtual void PostWalk(AnnotatedAssignStatement node) { }
+
         // AssertStatement
         public virtual bool Walk(AssertStatement node) { return true; }
         public virtual void PostWalk(AssertStatement node) { }
@@ -379,6 +383,10 @@ public override void PostWalk(YieldExpression node) { }
         public override bool Walk(YieldFromExpression node) { return false; }
         public override void PostWalk(YieldFromExpression node) { }
 
+        // AnnotatedAssignStatement
+        public override bool Walk(AnnotatedAssignStatement node) { return false; }
+        public override void PostWalk(AnnotatedAssignStatement node) { }
+
         // AssertStatement
         public override bool Walk(AssertStatement node) { return false; }
         public override void PostWalk(AssertStatement node) { }

Src/IronPython/Compiler/Parser.cs

@@ -660,14 +660,33 @@ private Statement FinishAssignments(Expression right) {
             return assign;
         }
 
-        // expr_stmt: testlist_star_expr (augassign (yield_expr|testlist) | ('=' (yield_expr|testlist_star_expr))*)
+        // expr_stmt: testlist_star_expr (annassign | augassign (yield_expr|testlist) | ('=' (yield_expr|testlist_star_expr))*)
         // augassign: ('+=' | '-=' | '*=' | '@=' | '/=' | '%=' | '&=' | '|=' | '^=' | '<<=' | '>>=' | '**=' | '//=')
         private Statement ParseExprStmt() {
             Expression ret = ParseTestListStarExpr();
             if (ret is ErrorExpression) {
                 NextToken();
             }
 
+            // annassign
+            if (MaybeEat(TokenKind.Colon)) {
+                if (ret.CheckAssign() is { } assignError) {
+                    ReportSyntaxError(ret.StartIndex, ret.EndIndex, assignError, ErrorCodes.SyntaxError | ErrorCodes.NoCaret);
+                }
+
+                var annotation = ParseTest();
+                if (MaybeEat(TokenKind.Assign)) {
+                    var right = ParseTest();
+                    var assign = new AnnotatedAssignStatement(ret, annotation, right, ret is NameExpression);
+                    assign.SetLoc(_globalParent, ret.StartIndex, right.EndIndex);
+                    return assign;
+                } else {
+                    var assign = new AnnotatedAssignStatement(ret, annotation, null, ret is NameExpression);
+                    assign.SetLoc(_globalParent, ret.IndexSpan);
+                    return assign;
+                }
+            }
+
             if (PeekToken(TokenKind.Assign)) {
                 return FinishAssignments(ret);
             }

Src/IronPython/Modules/_ast.cs

@@ -214,6 +214,7 @@ internal static stmt Convert(Statement stmt) {
                 stmt ast = stmt switch {
                     FunctionDefinition s => new FunctionDef(s),
                     ReturnStatement s => new Return(s),
+                    AnnotatedAssignStatement s => new AnnAssign(s),
                     AssignmentStatement s => new Assign(s),
                     AugmentedAssignStatement s => new AugAssign(s),
                     DelStatement s => new Delete(s),
@@ -711,6 +712,43 @@ public class And : boolop {
             internal static readonly And Instance = new And();
         }
 
+        [PythonType]
+        public class AnnAssign : stmt {
+            public AnnAssign() {
+                _fields = PythonTuple.MakeTuple(new[] { nameof(target), nameof(annotation), nameof(value), nameof(simple) });
+            }
+
+            public AnnAssign(expr target, expr annotation, expr value, int simple, [Optional] int? lineno, [Optional] int? col_offset)
+                : this() {
+                this.target = target;
+                this.annotation = annotation;
+                this.value = value;
+                this.simple = simple;
+                _lineno = lineno;
+                _col_offset = col_offset;
+            }
+
+            internal AnnAssign(AnnotatedAssignStatement stmt)
+                : this() {
+                target = Convert(stmt.Target, Store.Instance);
+                annotation = Convert(stmt.Annotation);
+                value = stmt.Value is null ? null : Convert(stmt.Value);
+                simple = stmt.Simple ? 1 : 0;
+            }
+
+            internal override Statement Revert() {
+                return new AnnotatedAssignStatement(expr.Revert(target), expr.Revert(annotation), expr.Revert(value), simple != 0);
+            }
+
+            public expr target { get; set; }
+
+            public expr annotation { get; set; }
+
+            public expr value { get; set; }
+
+            public int simple { get; set; }
+        }
+
         [PythonType]
         public class Assert : stmt {
             public Assert() {

Src/Scripts/generate_walker.py

@@ -31,7 +31,7 @@ def get_ast(assembly, roots):
         if node.IsAbstract: continue
 
         for root in roots:
-            if inherits(node, root) and node.IsPublic: 
+            if inherits(node, root) and node.IsPublic:
                 sets[root].add(node.Name)
                 break
 
@@ -40,7 +40,7 @@ def get_ast(assembly, roots):
         result.extend(sorted(list(sets[root])))
 
     return result
-    
+
 def gen_walker(cw, nodes, method, value):
     space = 0
     for node in nodes:
@@ -77,6 +77,7 @@ def gen_python_name_binder(cw):
                   "DelStatement",
                   "ClassDefinition",
                   "FunctionDefinition",
+                  "AnnotatedAssignStatement",
                   "AugmentedAssignStatement",
                   "AssignmentStatement",
                   "RaiseStatement",