Issue #27985: Implement PEP 526 -- Syntax for Variable Annotations.

Patch by Ivan Levkivskyi.

1 files changed, 123 insertions, 1 deletions

diff --git a/Python/Python-ast.c b/Python/Python-ast.c
index 1193c7c..6ab57df 100644
--- a/Python/Python-ast.c
+++ b/Python/Python-ast.c
@@ -86,6 +86,15 @@ static char *AugAssign_fields[]={
     "op",
     "value",
 };
+static PyTypeObject *AnnAssign_type;
+_Py_IDENTIFIER(annotation);
+_Py_IDENTIFIER(simple);
+static char *AnnAssign_fields[]={
+    "target",
+    "annotation",
+    "value",
+    "simple",
+};
 static PyTypeObject *For_type;
 _Py_IDENTIFIER(iter);
 _Py_IDENTIFIER(orelse);
@@ -466,7 +475,6 @@ static char *arg_attributes[] = {
     "col_offset",
 };
 _Py_IDENTIFIER(arg);
-_Py_IDENTIFIER(annotation);
 static char *arg_fields[]={
     "arg",
     "annotation",
@@ -873,6 +881,8 @@ static int init_types(void)
     if (!Assign_type) return 0;
     AugAssign_type = make_type("AugAssign", stmt_type, AugAssign_fields, 3);
     if (!AugAssign_type) return 0;
+    AnnAssign_type = make_type("AnnAssign", stmt_type, AnnAssign_fields, 4);
+    if (!AnnAssign_type) return 0;
     For_type = make_type("For", stmt_type, For_fields, 4);
     if (!For_type) return 0;
     AsyncFor_type = make_type("AsyncFor", stmt_type, AsyncFor_fields, 4);
@@ -1407,6 +1417,34 @@ AugAssign(expr_ty target, operator_ty op, expr_ty value, int lineno, int
 }
 
 stmt_ty
+AnnAssign(expr_ty target, expr_ty annotation, expr_ty value, int simple, int
+          lineno, int col_offset, PyArena *arena)
+{
+    stmt_ty p;
+    if (!target) {
+        PyErr_SetString(PyExc_ValueError,
+                        "field target is required for AnnAssign");
+        return NULL;
+    }
+    if (!annotation) {
+        PyErr_SetString(PyExc_ValueError,
+                        "field annotation is required for AnnAssign");
+        return NULL;
+    }
+    p = (stmt_ty)PyArena_Malloc(arena, sizeof(*p));
+    if (!p)
+        return NULL;
+    p->kind = AnnAssign_kind;
+    p->v.AnnAssign.target = target;
+    p->v.AnnAssign.annotation = annotation;
+    p->v.AnnAssign.value = value;
+    p->v.AnnAssign.simple = simple;
+    p->lineno = lineno;
+    p->col_offset = col_offset;
+    return p;
+}
+
+stmt_ty
 For(expr_ty target, expr_ty iter, asdl_seq * body, asdl_seq * orelse, int
     lineno, int col_offset, PyArena *arena)
 {
@@ -2740,6 +2778,30 @@ ast2obj_stmt(void* _o)
             goto failed;
         Py_DECREF(value);
         break;
+    case AnnAssign_kind:
+        result = PyType_GenericNew(AnnAssign_type, NULL, NULL);
+        if (!result) goto failed;
+        value = ast2obj_expr(o->v.AnnAssign.target);
+        if (!value) goto failed;
+        if (_PyObject_SetAttrId(result, &PyId_target, value) == -1)
+            goto failed;
+        Py_DECREF(value);
+        value = ast2obj_expr(o->v.AnnAssign.annotation);
+        if (!value) goto failed;
+        if (_PyObject_SetAttrId(result, &PyId_annotation, value) == -1)
+            goto failed;
+        Py_DECREF(value);
+        value = ast2obj_expr(o->v.AnnAssign.value);
+        if (!value) goto failed;
+        if (_PyObject_SetAttrId(result, &PyId_value, value) == -1)
+            goto failed;
+        Py_DECREF(value);
+        value = ast2obj_int(o->v.AnnAssign.simple);
+        if (!value) goto failed;
+        if (_PyObject_SetAttrId(result, &PyId_simple, value) == -1)
+            goto failed;
+        Py_DECREF(value);
+        break;
     case For_kind:
         result = PyType_GenericNew(For_type, NULL, NULL);
         if (!result) goto failed;
@@ -4535,6 +4597,64 @@ obj2ast_stmt(PyObject* obj, stmt_ty* out, PyArena* arena)
         if (*out == NULL) goto failed;
         return 0;
     }
+    isinstance = PyObject_IsInstance(obj, (PyObject*)AnnAssign_type);
+    if (isinstance == -1) {
+        return 1;
+    }
+    if (isinstance) {
+        expr_ty target;
+        expr_ty annotation;
+        expr_ty value;
+        int simple;
+
+        if (_PyObject_HasAttrId(obj, &PyId_target)) {
+            int res;
+            tmp = _PyObject_GetAttrId(obj, &PyId_target);
+            if (tmp == NULL) goto failed;
+            res = obj2ast_expr(tmp, &target, arena);
+            if (res != 0) goto failed;
+            Py_CLEAR(tmp);
+        } else {
+            PyErr_SetString(PyExc_TypeError, "required field \"target\" missing from AnnAssign");
+            return 1;
+        }
+        if (_PyObject_HasAttrId(obj, &PyId_annotation)) {
+            int res;
+            tmp = _PyObject_GetAttrId(obj, &PyId_annotation);
+            if (tmp == NULL) goto failed;
+            res = obj2ast_expr(tmp, &annotation, arena);
+            if (res != 0) goto failed;
+            Py_CLEAR(tmp);
+        } else {
+            PyErr_SetString(PyExc_TypeError, "required field \"annotation\" missing from AnnAssign");
+            return 1;
+        }
+        if (exists_not_none(obj, &PyId_value)) {
+            int res;
+            tmp = _PyObject_GetAttrId(obj, &PyId_value);
+            if (tmp == NULL) goto failed;
+            res = obj2ast_expr(tmp, &value, arena);
+            if (res != 0) goto failed;
+            Py_CLEAR(tmp);
+        } else {
+            value = NULL;
+        }
+        if (_PyObject_HasAttrId(obj, &PyId_simple)) {
+            int res;
+            tmp = _PyObject_GetAttrId(obj, &PyId_simple);
+            if (tmp == NULL) goto failed;
+            res = obj2ast_int(tmp, &simple, arena);
+            if (res != 0) goto failed;
+            Py_CLEAR(tmp);
+        } else {
+            PyErr_SetString(PyExc_TypeError, "required field \"simple\" missing from AnnAssign");
+            return 1;
+        }
+        *out = AnnAssign(target, annotation, value, simple, lineno, col_offset,
+                         arena);
+        if (*out == NULL) goto failed;
+        return 0;
+    }
     isinstance = PyObject_IsInstance(obj, (PyObject*)For_type);
     if (isinstance == -1) {
         return 1;
@@ -7517,6 +7637,8 @@ PyInit__ast(void)
         NULL;
     if (PyDict_SetItemString(d, "AugAssign", (PyObject*)AugAssign_type) < 0)
         return NULL;
+    if (PyDict_SetItemString(d, "AnnAssign", (PyObject*)AnnAssign_type) < 0)
+        return NULL;
     if (PyDict_SetItemString(d, "For", (PyObject*)For_type) < 0) return NULL;
     if (PyDict_SetItemString(d, "AsyncFor", (PyObject*)AsyncFor_type) < 0)
         return NULL;