bpo-36878: Track extra text added to 'type: ignore' in the AST (GH-13479)

GH-13238 made extra text after a # type: ignore accepted by the parser.
This finishes the job and actually plumbs the extra text through the
parser and makes it available in the AST.

2 files changed, 34 insertions, 4 deletions

diff --git a/Python/Python-ast.c b/Python/Python-ast.c
index 5527505..e84a758 100644
--- a/Python/Python-ast.c
+++ b/Python/Python-ast.c
@@ -524,8 +524,10 @@ static char *withitem_fields[]={
 static PyTypeObject *type_ignore_type;
 static PyObject* ast2obj_type_ignore(void*);
 static PyTypeObject *TypeIgnore_type;
+_Py_IDENTIFIER(tag);
 static char *TypeIgnore_fields[]={
     "lineno",
+    "tag",
 };
 
 
@@ -1164,7 +1166,7 @@ static int init_types(void)
     if (!type_ignore_type) return 0;
     if (!add_attributes(type_ignore_type, NULL, 0)) return 0;
     TypeIgnore_type = make_type("TypeIgnore", type_ignore_type,
-                                TypeIgnore_fields, 1);
+                                TypeIgnore_fields, 2);
     if (!TypeIgnore_type) return 0;
     initialized = 1;
     return 1;
@@ -2667,14 +2669,20 @@ withitem(expr_ty context_expr, expr_ty optional_vars, PyArena *arena)
 }
 
 type_ignore_ty
-TypeIgnore(int lineno, PyArena *arena)
+TypeIgnore(int lineno, string tag, PyArena *arena)
 {
     type_ignore_ty p;
+    if (!tag) {
+        PyErr_SetString(PyExc_ValueError,
+                        "field tag is required for TypeIgnore");
+        return NULL;
+    }
     p = (type_ignore_ty)PyArena_Malloc(arena, sizeof(*p));
     if (!p)
         return NULL;
     p->kind = TypeIgnore_kind;
     p->v.TypeIgnore.lineno = lineno;
+    p->v.TypeIgnore.tag = tag;
     return p;
 }
 
@@ -4158,6 +4166,11 @@ ast2obj_type_ignore(void* _o)
         if (_PyObject_SetAttrId(result, &PyId_lineno, value) == -1)
             goto failed;
         Py_DECREF(value);
+        value = ast2obj_string(o->v.TypeIgnore.tag);
+        if (!value) goto failed;
+        if (_PyObject_SetAttrId(result, &PyId_tag, value) == -1)
+            goto failed;
+        Py_DECREF(value);
         break;
     }
     return result;
@@ -8738,6 +8751,7 @@ obj2ast_type_ignore(PyObject* obj, type_ignore_ty* out, PyArena* arena)
     }
     if (isinstance) {
         int lineno;
+        string tag;
 
         if (_PyObject_LookupAttrId(obj, &PyId_lineno, &tmp) < 0) {
             return 1;
@@ -8752,7 +8766,20 @@ obj2ast_type_ignore(PyObject* obj, type_ignore_ty* out, PyArena* arena)
             if (res != 0) goto failed;
             Py_CLEAR(tmp);
         }
-        *out = TypeIgnore(lineno, arena);
+        if (_PyObject_LookupAttrId(obj, &PyId_tag, &tmp) < 0) {
+            return 1;
+        }
+        if (tmp == NULL) {
+            PyErr_SetString(PyExc_TypeError, "required field \"tag\" missing from TypeIgnore");
+            return 1;
+        }
+        else {
+            int res;
+            res = obj2ast_string(tmp, &tag, arena);
+            if (res != 0) goto failed;
+            Py_CLEAR(tmp);
+        }
+        *out = TypeIgnore(lineno, tag, arena);
         if (*out == NULL) goto failed;
         return 0;
     }
diff --git a/Python/ast.c b/Python/ast.c
index abc8d89..6259827 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -830,7 +830,10 @@ PyAST_FromNodeObject(const node *n, PyCompilerFlags *flags,
                 goto out;
 
             for (i = 0; i < num; i++) {
-                type_ignore_ty ti = TypeIgnore(LINENO(CHILD(ch, i)), arena);
+                string type_comment = new_type_comment(STR(CHILD(ch, i)), &c);
+                if (!type_comment)
+                    goto out;
+                type_ignore_ty ti = TypeIgnore(LINENO(CHILD(ch, i)), type_comment, arena);
                 if (!ti)
                    goto out;
                asdl_seq_SET(type_ignores, i, ti);