gh-122245: move checks for writes and shadowing of __debug__ to symtable (#122246)

2 files changed, 85 insertions, 83 deletions

diff --git a/Python/compile.c b/Python/compile.c
index 9707759..d07a435 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1954,55 +1954,6 @@ compiler_default_arguments(struct compiler *c, location loc,
     return funcflags;
 }
 
-static bool
-forbidden_name(struct compiler *c, location loc, identifier name,
-               expr_context_ty ctx)
-{
-    if (ctx == Store && _PyUnicode_EqualToASCIIString(name, "__debug__")) {
-        compiler_error(c, loc, "cannot assign to __debug__");
-        return true;
-    }
-    if (ctx == Del && _PyUnicode_EqualToASCIIString(name, "__debug__")) {
-        compiler_error(c, loc, "cannot delete __debug__");
-        return true;
-    }
-    return false;
-}
-
-static int
-compiler_check_debug_one_arg(struct compiler *c, arg_ty arg)
-{
-    if (arg != NULL) {
-        if (forbidden_name(c, LOC(arg), arg->arg, Store)) {
-            return ERROR;
-        }
-    }
-    return SUCCESS;
-}
-
-static int
-compiler_check_debug_args_seq(struct compiler *c, asdl_arg_seq *args)
-{
-    if (args != NULL) {
-        for (Py_ssize_t i = 0, n = asdl_seq_LEN(args); i < n; i++) {
-            RETURN_IF_ERROR(
-                compiler_check_debug_one_arg(c, asdl_seq_GET(args, i)));
-        }
-    }
-    return SUCCESS;
-}
-
-static int
-compiler_check_debug_args(struct compiler *c, arguments_ty args)
-{
-    RETURN_IF_ERROR(compiler_check_debug_args_seq(c, args->posonlyargs));
-    RETURN_IF_ERROR(compiler_check_debug_args_seq(c, args->args));
-    RETURN_IF_ERROR(compiler_check_debug_one_arg(c, args->vararg));
-    RETURN_IF_ERROR(compiler_check_debug_args_seq(c, args->kwonlyargs));
-    RETURN_IF_ERROR(compiler_check_debug_one_arg(c, args->kwarg));
-    return SUCCESS;
-}
-
 static int
 wrap_in_stopiteration_handler(struct compiler *c)
 {
@@ -2267,7 +2218,6 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
         type_params = s->v.FunctionDef.type_params;
     }
 
-    RETURN_IF_ERROR(compiler_check_debug_args(c, args));
     RETURN_IF_ERROR(compiler_decorators(c, decos));
 
     firstlineno = s->lineno;
@@ -2910,8 +2860,6 @@ compiler_lambda(struct compiler *c, expr_ty e)
     arguments_ty args = e->v.Lambda.args;
     assert(e->kind == Lambda_kind);
 
-    RETURN_IF_ERROR(compiler_check_debug_args(c, args));
-
     location loc = LOC(e);
     funcflags = compiler_default_arguments(c, loc, args);
     if (funcflags == -1) {
@@ -4086,10 +4034,6 @@ compiler_nameop(struct compiler *c, location loc,
            !_PyUnicode_EqualToASCIIString(name, "True") &&
            !_PyUnicode_EqualToASCIIString(name, "False"));
 
-    if (forbidden_name(c, loc, name, ctx)) {
-        return ERROR;
-    }
-
     mangled = compiler_maybe_mangle(c, name);
     if (!mangled) {
         return ERROR;
@@ -4878,10 +4822,6 @@ validate_keywords(struct compiler *c, asdl_keyword_seq *keywords)
         if (key->arg == NULL) {
             continue;
         }
-        location loc = LOC(key);
-        if (forbidden_name(c, loc, key->arg, Store)) {
-            return ERROR;
-        }
         for (Py_ssize_t j = i + 1; j < nkeywords; j++) {
             keyword_ty other = ((keyword_ty)asdl_seq_GET(keywords, j));
             if (other->arg && !PyUnicode_Compare(key->arg, other->arg)) {
@@ -6135,9 +6075,6 @@ compiler_visit_expr(struct compiler *c, expr_ty e)
             ADDOP_NAME(c, loc, LOAD_ATTR, e->v.Attribute.attr, names);
             break;
         case Store:
-            if (forbidden_name(c, loc, e->v.Attribute.attr, e->v.Attribute.ctx)) {
-                return ERROR;
-            }
             ADDOP_NAME(c, loc, STORE_ATTR, e->v.Attribute.attr, names);
             break;
         case Del:
@@ -6331,9 +6268,6 @@ compiler_annassign(struct compiler *c, stmt_ty s)
     }
     switch (targ->kind) {
     case Name_kind:
-        if (forbidden_name(c, loc, targ->v.Name.id, Store)) {
-            return ERROR;
-        }
         /* If we have a simple name in a module or class, store annotation. */
         if (s->v.AnnAssign.simple &&
             (c->u->u_scope_type == COMPILER_SCOPE_MODULE ||
@@ -6365,9 +6299,6 @@ compiler_annassign(struct compiler *c, stmt_ty s)
         }
         break;
     case Attribute_kind:
-        if (forbidden_name(c, loc, targ->v.Attribute.attr, Store)) {
-            return ERROR;
-        }
         if (!s->v.AnnAssign.value &&
             check_ann_expr(c, targ->v.Attribute.value) < 0) {
             return ERROR;
@@ -6631,9 +6562,6 @@ pattern_helper_store_name(struct compiler *c, location loc,
         ADDOP(c, loc, POP_TOP);
         return SUCCESS;
     }
-    if (forbidden_name(c, loc, n, Store)) {
-        return ERROR;
-    }
     // Can't assign to the same name twice:
     int duplicate = PySequence_Contains(pc->stores, n);
     RETURN_IF_ERROR(duplicate);
@@ -6791,10 +6719,6 @@ validate_kwd_attrs(struct compiler *c, asdl_identifier_seq *attrs, asdl_pattern_
     Py_ssize_t nattrs = asdl_seq_LEN(attrs);
     for (Py_ssize_t i = 0; i < nattrs; i++) {
         identifier attr = ((identifier)asdl_seq_GET(attrs, i));
-        location loc = LOC((pattern_ty) asdl_seq_GET(patterns, i));
-        if (forbidden_name(c, loc, attr, Store)) {
-            return ERROR;
-        }
         for (Py_ssize_t j = i + 1; j < nattrs; j++) {
             identifier other = ((identifier)asdl_seq_GET(attrs, j));
             if (!PyUnicode_Compare(attr, other)) {
diff --git a/Python/symtable.c b/Python/symtable.c
index c4508ca..a5fa758 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -1495,8 +1495,57 @@ error:
 }
 
 static int
-symtable_add_def(struct symtable *st, PyObject *name, int flag, _Py_SourceLocation loc)
+check_name(struct symtable *st, PyObject *name, _Py_SourceLocation loc,
+           expr_context_ty ctx)
 {
+    if (ctx == Store && _PyUnicode_EqualToASCIIString(name, "__debug__")) {
+        PyErr_SetString(PyExc_SyntaxError, "cannot assign to __debug__");
+        SET_ERROR_LOCATION(st->st_filename, loc);
+        return 0;
+    }
+    if (ctx == Del && _PyUnicode_EqualToASCIIString(name, "__debug__")) {
+        PyErr_SetString(PyExc_SyntaxError, "cannot delete __debug__");
+        SET_ERROR_LOCATION(st->st_filename, loc);
+        return 0;
+    }
+    return 1;
+}
+
+static int
+check_keywords(struct symtable *st, asdl_keyword_seq *keywords)
+{
+    for (Py_ssize_t i = 0; i < asdl_seq_LEN(keywords); i++) {
+        keyword_ty key = ((keyword_ty)asdl_seq_GET(keywords, i));
+        if (key->arg  && !check_name(st, key->arg, LOCATION(key), Store)) {
+            return 0;
+        }
+    }
+    return 1;
+}
+
+static int
+check_kwd_patterns(struct symtable *st, pattern_ty p)
+{
+    assert(p->kind == MatchClass_kind);
+    asdl_identifier_seq *kwd_attrs = p->v.MatchClass.kwd_attrs;
+    asdl_pattern_seq *kwd_patterns = p->v.MatchClass.kwd_patterns;
+    for (Py_ssize_t i = 0; i < asdl_seq_LEN(kwd_attrs); i++) {
+        _Py_SourceLocation loc = LOCATION(asdl_seq_GET(kwd_patterns, i));
+        if (!check_name(st, asdl_seq_GET(kwd_attrs, i), loc, Store)) {
+            return 0;
+        }
+    }
+    return 1;
+}
+
+static int
+symtable_add_def_ctx(struct symtable *st, PyObject *name, int flag,
+                     _Py_SourceLocation loc, expr_context_ty ctx)
+{
+    int write_mask = DEF_PARAM | DEF_LOCAL | DEF_IMPORT;
+    if ((flag & write_mask) && !check_name(st, name, loc, ctx)) {
+        return 0;
+    }
     if ((flag & DEF_TYPE_PARAM) && st->st_cur->ste_mangled_names != NULL) {
         if(PySet_Add(st->st_cur->ste_mangled_names, name) < 0) {
             return 0;
@@ -1506,6 +1555,14 @@ symtable_add_def(struct symtable *st, PyObject *name, int flag, _Py_SourceLocati
 }
 
 static int
+symtable_add_def(struct symtable *st, PyObject *name, int flag,
+                 _Py_SourceLocation loc)
+{
+    return symtable_add_def_ctx(st, name, flag, loc,
+                                flag == USE ? Load : Store);
+}
+
+static int
 symtable_enter_type_param_block(struct symtable *st, identifier name,
                                void *ast, int has_defaults, int has_kwdefaults,
                                enum _stmt_kind kind, _Py_SourceLocation loc)
@@ -1757,6 +1814,9 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             VISIT_SEQ(st, type_param, s->v.ClassDef.type_params);
         }
         VISIT_SEQ(st, expr, s->v.ClassDef.bases);
+        if (!check_keywords(st, s->v.ClassDef.keywords)) {
+            VISIT_QUIT(st, 0);
+        }
         VISIT_SEQ(st, keyword, s->v.ClassDef.keywords);
         if (!symtable_enter_block(st, s->v.ClassDef.name, ClassBlock,
                                   (void *)s, LOCATION(s))) {
@@ -1871,10 +1931,11 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             VISIT(st, expr, s->v.AnnAssign.value);
         }
         break;
-    case AugAssign_kind:
+    case AugAssign_kind: {
         VISIT(st, expr, s->v.AugAssign.target);
         VISIT(st, expr, s->v.AugAssign.value);
         break;
+    }
     case For_kind:
         VISIT(st, expr, s->v.For.target);
         VISIT(st, expr, s->v.For.iter);
@@ -2311,6 +2372,9 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
     case Call_kind:
         VISIT(st, expr, e->v.Call.func);
         VISIT_SEQ(st, expr, e->v.Call.args);
+        if (!check_keywords(st, e->v.Call.keywords)) {
+            VISIT_QUIT(st, 0);
+        }
         VISIT_SEQ_WITH_NULL(st, keyword, e->v.Call.keywords);
         break;
     case FormattedValue_kind:
@@ -2326,6 +2390,9 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
         break;
     /* The following exprs can be assignment targets. */
     case Attribute_kind:
+        if (!check_name(st, e->v.Attribute.attr, LOCATION(e), e->v.Attribute.ctx)) {
+            VISIT_QUIT(st, 0);
+        }
         VISIT(st, expr, e->v.Attribute.value);
         break;
     case Subscript_kind:
@@ -2344,9 +2411,11 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
             VISIT(st, expr, e->v.Slice.step);
         break;
     case Name_kind:
-        if (!symtable_add_def(st, e->v.Name.id,
-                              e->v.Name.ctx == Load ? USE : DEF_LOCAL, LOCATION(e)))
+        if (!symtable_add_def_ctx(st, e->v.Name.id,
+                                  e->v.Name.ctx == Load ? USE : DEF_LOCAL,
+                                  LOCATION(e), e->v.Name.ctx)) {
             VISIT_QUIT(st, 0);
+        }
         /* Special-case super: it counts as a use of __class__ */
         if (e->v.Name.ctx == Load &&
             _PyST_IsFunctionLike(st->st_cur) &&
@@ -2472,19 +2541,26 @@ symtable_visit_pattern(struct symtable *st, pattern_ty p)
         break;
     case MatchStar_kind:
         if (p->v.MatchStar.name) {
-            symtable_add_def(st, p->v.MatchStar.name, DEF_LOCAL, LOCATION(p));
+            if (!symtable_add_def(st, p->v.MatchStar.name, DEF_LOCAL, LOCATION(p))) {
+                VISIT_QUIT(st, 0);
+            }
         }
         break;
     case MatchMapping_kind:
         VISIT_SEQ(st, expr, p->v.MatchMapping.keys);
         VISIT_SEQ(st, pattern, p->v.MatchMapping.patterns);
         if (p->v.MatchMapping.rest) {
-            symtable_add_def(st, p->v.MatchMapping.rest, DEF_LOCAL, LOCATION(p));
+            if (!symtable_add_def(st, p->v.MatchMapping.rest, DEF_LOCAL, LOCATION(p))) {
+                VISIT_QUIT(st, 0);
+            }
         }
         break;
     case MatchClass_kind:
         VISIT(st, expr, p->v.MatchClass.cls);
         VISIT_SEQ(st, pattern, p->v.MatchClass.patterns);
+        if (!check_kwd_patterns(st, p)) {
+            VISIT_QUIT(st, 0);
+        }
         VISIT_SEQ(st, pattern, p->v.MatchClass.kwd_patterns);
         break;
     case MatchAs_kind:
@@ -2492,7 +2568,9 @@ symtable_visit_pattern(struct symtable *st, pattern_ty p)
             VISIT(st, pattern, p->v.MatchAs.pattern);
         }
         if (p->v.MatchAs.name) {
-            symtable_add_def(st, p->v.MatchAs.name, DEF_LOCAL, LOCATION(p));
+            if (!symtable_add_def(st, p->v.MatchAs.name, DEF_LOCAL, LOCATION(p))) {
+                VISIT_QUIT(st, 0);
+            }
         }
         break;
     case MatchOr_kind: