Check for NULL from NEW_TYPE_COMMENT()

Also:
- remove redundant 'c' argument from new_type_comment()
- double-check that TYPE_COMMENT in a suite is followed by NL

diff --git a/Python/ast.c b/Python/ast.c
index cad82d1d582777cd70f1b27ec4e6458790340511..2d3bfc803057699caf6bd73b6e0c47a6d7dcf1b0 100644 (file)
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -699,11 +699,11 @@ ast_error(struct compiling *c, const node *n, const char *errmsg, ...)
 */
 
 static string
-new_type_comment(const char *s, struct compiling *c)
+new_type_comment(const char *s)
 {
   return PyUnicode_DecodeUTF8(s, strlen(s), NULL);
 }
-#define NEW_TYPE_COMMENT(n) new_type_comment(STR(n), c)
+#define NEW_TYPE_COMMENT(n) new_type_comment(STR(n))
 
 static int
 num_stmts(const node *n)
@@ -1412,6 +1412,8 @@ handle_keywordonly_args(struct compiling *c, const node *n, int start,
             case TYPE_COMMENT:
                 /* arg will be equal to the last argument processed */
                 arg->type_comment = NEW_TYPE_COMMENT(ch);
+                if (!arg->type_comment)
+                    goto error;
                 i += 1;
                 break;
             case DOUBLESTAR:
@@ -1581,6 +1583,8 @@ ast_for_arguments(struct compiling *c, const node *n)
 
                 if (i < NCH(n) && TYPE(CHILD(n, i)) == TYPE_COMMENT) {
                         vararg->type_comment = NEW_TYPE_COMMENT(CHILD(n, i));
+                        if (!vararg->type_comment)
+                            return NULL;
                         i += 1;
                     }
 
@@ -1612,6 +1616,8 @@ ast_for_arguments(struct compiling *c, const node *n)
 
                 /* arg will be equal to the last argument processed */
                 arg->type_comment = NEW_TYPE_COMMENT(ch);
+                if (!arg->type_comment)
+                    return NULL;
                 i += 1;
                 break;
             default:
@@ -1750,6 +1756,8 @@ ast_for_funcdef_impl(struct compiling *c, const node *n0,
     }
     if (TYPE(CHILD(n, name_i + 3)) == TYPE_COMMENT) {
         type_comment = NEW_TYPE_COMMENT(CHILD(n, name_i + 3));
+        if (!type_comment)
+            return NULL;
         name_i += 1;
     }
     body = ast_for_suite(c, CHILD(n, name_i + 3));
@@ -1767,6 +1775,8 @@ ast_for_funcdef_impl(struct compiling *c, const node *n0,
                 return NULL;
             }
             type_comment = NEW_TYPE_COMMENT(tc);
+            if (!type_comment)
+                return NULL;
         }
     }
 
@@ -3345,8 +3355,11 @@ ast_for_expr_stmt(struct compiling *c, const node *n)
             expression = ast_for_expr(c, value);
         if (!expression)
             return NULL;
-        if (has_type_comment)
+        if (has_type_comment) {
             type_comment = NEW_TYPE_COMMENT(CHILD(n, nch_minus_type));
+            if (!type_comment)
+                return NULL;
+        }
         else
             type_comment = NULL;
         return Assign(targets, expression, type_comment, LINENO(n), n->n_col_offset,
@@ -3809,8 +3822,10 @@ ast_for_suite(struct compiling *c, const node *n)
     }
     else {
         i = 2;
-        if (TYPE(CHILD(n, 1)) == TYPE_COMMENT)
+        if (TYPE(CHILD(n, 1)) == TYPE_COMMENT) {
             i += 2;
+            REQ(CHILD(n, 2), NEWLINE);
+        }
 
         for (; i < (NCH(n) - 1); i++) {
             ch = CHILD(n, i);
@@ -4086,8 +4101,11 @@ ast_for_for_stmt(struct compiling *c, const node *n0, bool is_async)
         get_last_end_pos(suite_seq, &end_lineno, &end_col_offset);
     }
 
-    if (has_type_comment)
+    if (has_type_comment) {
         type_comment = NEW_TYPE_COMMENT(CHILD(n, 5));
+        if (!type_comment)
+            return NULL;
+    }
     else
         type_comment = NULL;
 
@@ -4294,8 +4312,11 @@ ast_for_with_stmt(struct compiling *c, const node *n0, bool is_async)
         return NULL;
     get_last_end_pos(body, &end_lineno, &end_col_offset);
 
-    if (has_type_comment)
+    if (has_type_comment) {
         type_comment = NEW_TYPE_COMMENT(CHILD(n, NCH(n) - 2));
+        if (!type_comment)
+            return NULL;
+    }
     else
         type_comment = NULL;