bpo-36280: Add Constant.kind field (GH-12295)

The value is a string for string and byte literals, None otherwise.
It is 'u' for u"..." literals, 'b' for b"..." literals, '' for "..." literals.
The 'r' (raw) prefix is ignored.
Does not apply to f-strings.

This appears sufficient to make mypy capable of using the stdlib ast module instead of typed_ast (assuming a mypy patch I'm working on).

WIP: I need to make the tests pass. @ilevkivskyi @serhiy-storchaka 



https://bugs.python.org/issue36280

1 files changed, 46 insertions, 7 deletions

diff --git a/Python/ast.c b/Python/ast.c
index d4ee1b3..971b8dd 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -2313,13 +2313,13 @@ ast_for_atom(struct compiling *c, const node *n)
         size_t len = strlen(s);
         if (len >= 4 && len <= 5) {
             if (!strcmp(s, "None"))
-                return Constant(Py_None, LINENO(n), n->n_col_offset,
+                return Constant(Py_None, NULL, LINENO(n), n->n_col_offset,
                                 n->n_end_lineno, n->n_end_col_offset, c->c_arena);
             if (!strcmp(s, "True"))
-                return Constant(Py_True, LINENO(n), n->n_col_offset,
+                return Constant(Py_True, NULL, LINENO(n), n->n_col_offset,
                                 n->n_end_lineno, n->n_end_col_offset, c->c_arena);
             if (!strcmp(s, "False"))
-                return Constant(Py_False, LINENO(n), n->n_col_offset,
+                return Constant(Py_False, NULL, LINENO(n), n->n_col_offset,
                                 n->n_end_lineno, n->n_end_col_offset, c->c_arena);
         }
         name = new_identifier(s, c);
@@ -2374,11 +2374,11 @@ ast_for_atom(struct compiling *c, const node *n)
             Py_DECREF(pynum);
             return NULL;
         }
-        return Constant(pynum, LINENO(n), n->n_col_offset,
+        return Constant(pynum, NULL, LINENO(n), n->n_col_offset,
                         n->n_end_lineno, n->n_end_col_offset, c->c_arena);
     }
     case ELLIPSIS: /* Ellipsis */
-        return Constant(Py_Ellipsis, LINENO(n), n->n_col_offset,
+        return Constant(Py_Ellipsis, NULL, LINENO(n), n->n_col_offset,
                         n->n_end_lineno, n->n_end_col_offset, c->c_arena);
     case LPAR: /* some parenthesized expressions */
         ch = CHILD(n, 1);
@@ -5388,18 +5388,57 @@ FstringParser_Dealloc(FstringParser *state)
     ExprList_Dealloc(&state->expr_list);
 }
 
+/* Constants for the following */
+static PyObject *u_kind;
+
+/* Compute 'kind' field for string Constant (either 'u' or None) */
+static PyObject *
+make_kind(struct compiling *c, const node *n)
+{
+    char *s = NULL;
+    PyObject *kind = NULL;
+
+    /* Find the first string literal, if any */
+    while (TYPE(n) != STRING) {
+        if (NCH(n) == 0)
+            return NULL;
+        n = CHILD(n, 0);
+    }
+    REQ(n, STRING);
+
+    /* If it starts with 'u', return a PyUnicode "u" string */
+    s = STR(n);
+    if (s && *s == 'u') {
+        if (!u_kind) {
+            u_kind = PyUnicode_InternFromString("u");
+            if (!u_kind)
+                return NULL;
+        }
+        kind = u_kind;
+        if (PyArena_AddPyObject(c->c_arena, kind) < 0) {
+            return NULL;
+        }
+        Py_INCREF(kind);
+    }
+    return kind;
+}
+
 /* Make a Constant node, but decref the PyUnicode object being added. */
 static expr_ty
 make_str_node_and_del(PyObject **str, struct compiling *c, const node* n)
 {
     PyObject *s = *str;
+    PyObject *kind = NULL;
     *str = NULL;
     assert(PyUnicode_CheckExact(s));
     if (PyArena_AddPyObject(c->c_arena, s) < 0) {
         Py_DECREF(s);
         return NULL;
     }
-    return Constant(s, LINENO(n), n->n_col_offset,
+    kind = make_kind(c, n);
+    if (kind == NULL && PyErr_Occurred())
+        return NULL;
+    return Constant(s, kind, LINENO(n), n->n_col_offset,
                     n->n_end_lineno, n->n_end_col_offset, c->c_arena);
 }
 
@@ -5774,7 +5813,7 @@ parsestrplus(struct compiling *c, const node *n)
         /* Just return the bytes object and we're done. */
         if (PyArena_AddPyObject(c->c_arena, bytes_str) < 0)
             goto error;
-        return Constant(bytes_str, LINENO(n), n->n_col_offset,
+        return Constant(bytes_str, NULL, LINENO(n), n->n_col_offset,
                         n->n_end_lineno, n->n_end_col_offset, c->c_arena);
     }