[3.11] pythongh-96497: Mangle name before symtable lookup in 'symtable_extend_namedexpr_scope' (pythonGH-96561)

wookie184 · serhiy-storchaka · commit 38d66778b7b1 · 2024-02-17T14:53:18.000+02:00
(cherry picked from commit 664965a) Co-authored-by: wookie184 <wookie1840@gmail.com>
diff --git a/Lib/test/test_named_expressions.py b/Lib/test/test_named_expressions.py
@@ -222,6 +222,16 @@ def test_named_expression_invalid_set_comprehension_iterable_expression(self):
                 with self.assertRaisesRegex(SyntaxError, msg):
                     exec(f"lambda: {code}", {}) # Function scope
 
+    def test_named_expression_invalid_mangled_class_variables(self):
+        code = """class Foo:
+            def bar(self):
+                [[(__x:=2) for _ in range(2)] for __x in range(2)]
+        """
+
+        with self.assertRaisesRegex(SyntaxError,
+            "assignment expression cannot rebind comprehension iteration variable '__x'"):
+            exec(code, {}, {})
+
 
 class NamedExpressionAssignmentTest(unittest.TestCase):
 
@@ -598,6 +608,18 @@ def test_named_expression_scope_in_genexp(self):
         for idx, elem in enumerate(genexp):
             self.assertEqual(elem, b[idx] + a)
 
+    def test_named_expression_scope_mangled_names(self):
+        class Foo:
+            def f(self_):
+                global __x1
+                __x1 = 0
+                [_Foo__x1 := 1 for a in [2]]
+                self.assertEqual(__x1, 1)
+                [__x1 := 2 for a in [3]]
+                self.assertEqual(__x1, 2)
+
+        Foo().f()
+        self.assertEqual(_Foo__x1, 2)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-09-04-16-51-56.gh-issue-96497.HTBuIL.rst b/Misc/NEWS.d/next/Core and Builtins/2022-09-04-16-51-56.gh-issue-96497.HTBuIL.rst
@@ -0,0 +1,2 @@
+Fix incorrect resolution of mangled class variables used in assignment
+expressions in comprehensions.
diff --git a/Python/symtable.c b/Python/symtable.c
@@ -1042,16 +1042,22 @@ symtable_enter_block(struct symtable *st, identifier name, _Py_block_ty block,
 }
 
 static long
-symtable_lookup(struct symtable *st, PyObject *name)
+symtable_lookup_entry(struct symtable *st, PySTEntryObject *ste, PyObject *name)
 {
     PyObject *mangled = _Py_Mangle(st->st_private, name);
     if (!mangled)
         return 0;
-    long ret = _PyST_GetSymbol(st->st_cur, mangled);
+    long ret = _PyST_GetSymbol(ste, mangled);
     Py_DECREF(mangled);
     return ret;
 }
 
+static long
+symtable_lookup(struct symtable *st, PyObject *name)
+{
+    return symtable_lookup_entry(st, st->st_cur, name);
+}
+
 static int
 symtable_add_def_helper(struct symtable *st, PyObject *name, int flag, struct _symtable_entry *ste,
                         int lineno, int col_offset, int end_lineno, int end_col_offset)
@@ -1525,7 +1531,7 @@ symtable_extend_namedexpr_scope(struct symtable *st, expr_ty e)
          * binding conflict with iteration variables, otherwise skip it
          */
         if (ste->ste_comprehension) {
-            long target_in_scope = _PyST_GetSymbol(ste, target_name);
+            long target_in_scope = symtable_lookup_entry(st, ste, target_name);
             if (target_in_scope & DEF_COMP_ITER) {
                 PyErr_Format(PyExc_SyntaxError, NAMED_EXPR_COMP_CONFLICT, target_name);
                 PyErr_RangedSyntaxLocationObject(st->st_filename,
@@ -1540,7 +1546,7 @@ symtable_extend_namedexpr_scope(struct symtable *st, expr_ty e)
 
         /* If we find a FunctionBlock entry, add as GLOBAL/LOCAL or NONLOCAL/LOCAL */
         if (ste->ste_type == FunctionBlock) {
-            long target_in_scope = _PyST_GetSymbol(ste, target_name);
+            long target_in_scope = symtable_lookup_entry(st, ste, target_name);
             if (target_in_scope & DEF_GLOBAL) {
                 if (!symtable_add_def(st, target_name, DEF_GLOBAL, LOCATION(e)))
                     VISIT_QUIT(st, 0);

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+Fix incorrect resolution of mangled class variables used in assignment`
	`2`	`+expressions in comprehensions.`