pythongh-122707: move an error check from symtable to ast verification

Include/internal/pycore_symtable.h

@@ -123,7 +123,6 @@ typedef struct _symtable_entry {
     unsigned ste_comp_iter_target : 1; /* true if visiting comprehension target */
     unsigned ste_can_see_class_scope : 1; /* true if this block can see names bound in an
                                              enclosing class scope */
-    int ste_comp_iter_expr; /* non-zero if visiting a comprehension range expression */
     _Py_SourceLocation ste_loc; /* source location of block */
     struct _symtable_entry *ste_annotation_block; /* symbol table entry for this entry's annotations */
     struct symtable *ste_table;

Python/ast.c

@@ -12,6 +12,7 @@
 struct validator {
     int recursion_depth;            /* current recursion depth */
     int recursion_limit;            /* recursion limit */
+    int in_comp_iter_expr;          /* >0 in a comprehension iter expr */
 };
 
 #define ENTER_RECURSIVE(ST) \
@@ -38,6 +39,9 @@ static int validate_expr(struct validator *, expr_ty, expr_context_ty);
 static int validate_pattern(struct validator *, pattern_ty, int);
 static int validate_typeparam(struct validator *, type_param_ty);
 
+#define NAMED_EXPR_COMP_ITER_EXPR \
+"assignment expression cannot be used in a comprehension iterable expression"
+
 #define VALIDATE_POSITIONS(node) \
     if (node->lineno > node->end_lineno) { \
         PyErr_Format(PyExc_ValueError, \
@@ -89,10 +93,21 @@ validate_comprehension(struct validator *state, asdl_comprehension_seq *gens)
     }
     for (Py_ssize_t i = 0; i < asdl_seq_LEN(gens); i++) {
         comprehension_ty comp = asdl_seq_GET(gens, i);
-        if (!validate_expr(state, comp->target, Store) ||
-            !validate_expr(state, comp->iter, Load) ||
-            !validate_exprs(state, comp->ifs, Load, 0))
+
+        if (!validate_expr(state, comp->target, Store)) {
+            return 0;
+        }
+
+        state->in_comp_iter_expr++;
+        int res = validate_expr(state, comp->iter, Load);
+        state->in_comp_iter_expr--;
+        if (!res) {
             return 0;
+        }
+
+        if (!validate_exprs(state, comp->ifs, Load, 0)) {
+            return 0;
+        }
     }
     return 1;
 }
@@ -392,6 +407,10 @@ validate_expr(struct validator *state, expr_ty exp, expr_context_ty ctx)
                             "NamedExpr target must be a Name");
             return 0;
         }
+        if (state->in_comp_iter_expr) {
+            PyErr_SetString(PyExc_SyntaxError, NAMED_EXPR_COMP_ITER_EXPR);
+            return 0;
+        }
         ret = validate_expr(state, exp->v.NamedExpr.value, Load);
         break;
     /* This last case doesn't have any checking. */
@@ -1060,6 +1079,7 @@ _PyAST_Validate(mod_ty mod)
     starting_recursion_depth = recursion_depth;
     state.recursion_depth = starting_recursion_depth;
     state.recursion_limit = Py_C_RECURSION_LIMIT;
+    state.in_comp_iter_expr = 0;
 
     switch (mod->kind) {
     case Module_kind:

Python/symtable.c

@@ -52,9 +52,6 @@
 #define NAMED_EXPR_COMP_INNER_LOOP_CONFLICT \
 "comprehension inner loop cannot rebind assignment expression target '%U'"
 
-#define NAMED_EXPR_COMP_ITER_EXPR \
-"assignment expression cannot be used in a comprehension iterable expression"
-
 #define ANNOTATION_NOT_ALLOWED \
 "%s cannot be used within an annotation"
 
@@ -132,7 +129,6 @@ ste_new(struct symtable *st, identifier name, _Py_block_ty block,
     ste->ste_comp_inlined = 0;
     ste->ste_comp_iter_target = 0;
     ste->ste_can_see_class_scope = 0;
-    ste->ste_comp_iter_expr = 0;
     ste->ste_needs_classdict = 0;
     ste->ste_annotation_block = NULL;
 
@@ -1335,13 +1331,6 @@ symtable_enter_existing_block(struct symtable *st, PySTEntryObject* ste)
         return 0;
     }
     PySTEntryObject *prev = st->st_cur;
-    /* bpo-37757: For now, disallow *all* assignment expressions in the
-     * outermost iterator expression of a comprehension, even those inside
-     * a nested comprehension or a lambda expression.
-     */
-    if (prev) {
-        ste->ste_comp_iter_expr = prev->ste_comp_iter_expr;
-    }
     /* No need to inherit ste_mangled_names in classes, where all names
      * are mangled. */
     if (prev && prev->ste_mangled_names != NULL && ste->ste_type != ClassBlock) {
@@ -2249,12 +2238,6 @@ symtable_extend_namedexpr_scope(struct symtable *st, expr_ty e)
 static int
 symtable_handle_namedexpr(struct symtable *st, expr_ty e)
 {
-    if (st->st_cur->ste_comp_iter_expr > 0) {
-        /* Assignment isn't allowed in a comprehension iterable expression */
-        PyErr_Format(PyExc_SyntaxError, NAMED_EXPR_COMP_ITER_EXPR);
-        SET_ERROR_LOCATION(st->st_filename, LOCATION(e));
-        return 0;
-    }
     if (st->st_cur->ste_comprehension) {
         /* Inside a comprehension body, so find the right target scope */
         if (!symtable_extend_namedexpr_scope(st, e->v.NamedExpr.target))
@@ -2806,9 +2789,7 @@ symtable_visit_comprehension(struct symtable *st, comprehension_ty lc)
     st->st_cur->ste_comp_iter_target = 1;
     VISIT(st, expr, lc->target);
     st->st_cur->ste_comp_iter_target = 0;
-    st->st_cur->ste_comp_iter_expr++;
     VISIT(st, expr, lc->iter);
-    st->st_cur->ste_comp_iter_expr--;
     VISIT_SEQ(st, expr, lc->ifs);
     if (lc->is_async) {
         st->st_cur->ste_coroutine = 1;
@@ -2834,9 +2815,7 @@ symtable_handle_comprehension(struct symtable *st, expr_ty e,
     comprehension_ty outermost = ((comprehension_ty)
                                     asdl_seq_GET(generators, 0));
     /* Outermost iterator is evaluated in current scope */
-    st->st_cur->ste_comp_iter_expr++;
     VISIT(st, expr, outermost->iter);
-    st->st_cur->ste_comp_iter_expr--;
     /* Create comprehension scope for the rest */
     if (!scope_name ||
         !symtable_enter_block(st, scope_name, FunctionBlock, (void *)e, LOCATION(e))) {