pythongh-104602: ensure all cellvars are known up front

Include/internal/pycore_symtable.h

@@ -120,14 +120,15 @@ extern PyObject* _Py_Mangle(PyObject *p, PyObject *name);
 #define DEF_ANNOT 2<<7 /* this name is annotated */
 #define DEF_COMP_ITER 2<<8 /* this name is a comprehension iteration variable */
 #define DEF_TYPE_PARAM 2<<9 /* this name is a type parameter */
+#define DEF_COMP_CELL 2<<10 /* this name is a cell in an inlined comprehension */
 
 #define DEF_BOUND (DEF_LOCAL | DEF_PARAM | DEF_IMPORT)
 
 /* GLOBAL_EXPLICIT and GLOBAL_IMPLICIT are used internally by the symbol
  table. GLOBAL is returned from PyST_GetScope() for either of them.
- It is stored in ste_symbols at bits 12-15.
+ It is stored in ste_symbols at bits 13-16.
 */
-#define SCOPE_OFFSET 11
+#define SCOPE_OFFSET 12
 #define SCOPE_MASK (DEF_GLOBAL | DEF_LOCAL | DEF_PARAM | DEF_NONLOCAL)
 
 #define LOCAL 1

Lib/test/test_listcomps.py

@@ -379,6 +379,17 @@ def f():
  with self.assertRaises(UnboundLocalError):
  f()
 
+ def test_global_outside_cellvar_inside_plus_freevar(self):
+ code = """
+ a = 1
+ def f():
+ [(lambda: b) for b in [a]]
+ return b
+ x = f()
+ """
+ self._check_in_scopes(
+ code, {"x": 2}, ns={"b": 2}, scopes=["function", "module"])
+
 
 __test__ = {'doctests' : doctests}
 

Python/compile.c

@@ -1250,7 +1250,7 @@ compiler_enter_scope(struct compiler *c, identifier name,
  }
  u->u_metadata.u_name = Py_NewRef(name);
  u->u_metadata.u_varnames = list2dict(u->u_ste->ste_varnames);
- u->u_metadata.u_cellvars = dictbytype(u->u_ste->ste_symbols, CELL, 0, 0);
+ u->u_metadata.u_cellvars = dictbytype(u->u_ste->ste_symbols, CELL, DEF_COMP_CELL, 0);
  if (!u->u_metadata.u_varnames || !u->u_metadata.u_cellvars) {
  compiler_unit_free(u);
  return ERROR;

Python/symtable.c

@@ -632,7 +632,7 @@ is_free_in_any_child(PySTEntryObject *entry, PyObject *key)
 static int
 inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
  PyObject *scopes, PyObject *comp_free,
- PyObject *promote_to_cell)
+ PyObject *inlined_cells)
 {
  PyObject *k, *v;
  Py_ssize_t pos = 0;
@@ -645,6 +645,11 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
  }
  int scope = (comp_flags >> SCOPE_OFFSET) & SCOPE_MASK;
  int only_flags = comp_flags & ((1 << SCOPE_OFFSET) - 1);
+ if (scope == CELL) {
+ if (PySet_Add(inlined_cells, k) < 0) {
+ return 0;
+ }
+ }
  PyObject *existing = PyDict_GetItemWithError(ste->ste_symbols, k);
  if (existing == NULL && PyErr_Occurred()) {
  return 0;
@@ -665,14 +670,6 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
  }
  else {
  if (PyLong_AsLong(existing) & DEF_BOUND) {
- // cell vars in comprehension that are locals in outer scope
- // must be promoted to cell so u_cellvars isn't wrong
- if (scope == CELL && _PyST_IsFunctionLike(ste)) {
- if (PySet_Add(promote_to_cell, k) < 0) {
- return 0;
- }
- }
-
  // free vars in comprehension that are locals in outer scope can
  // now simply be locals, unless they are free in comp children
  if (!is_free_in_any_child(comp, k)) {
@@ -697,7 +694,7 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
 */
 
 static int
-analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
+analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells)
 {
  PyObject *name, *v, *v_cell;
  int success = 0;
@@ -712,7 +709,7 @@ analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
  scope = PyLong_AS_LONG(v);
  if (scope != LOCAL)
  continue;
- if (!PySet_Contains(free, name) && !PySet_Contains(promote_to_cell, name))
+ if (!PySet_Contains(free, name) && !PySet_Contains(inlined_cells, name))
  continue;
  /* Replace LOCAL with CELL for this name, and remove
  from free. It is safe to replace the value of name
@@ -752,7 +749,8 @@ drop_class_free(PySTEntryObject *ste, PyObject *free)
 */
 static int
 update_symbols(PyObject *symbols, PyObject *scopes,
- PyObject *bound, PyObject *free, int classflag)
+ PyObject *bound, PyObject *free,
+ PyObject *inlined_cells, int classflag)
 {
  PyObject *name = NULL, *itr = NULL;
  PyObject *v = NULL, *v_scope = NULL, *v_new = NULL, *v_free = NULL;
@@ -763,6 +761,9 @@ update_symbols(PyObject *symbols, PyObject *scopes,
  long scope, flags;
  assert(PyLong_Check(v));
  flags = PyLong_AS_LONG(v);
+ if (PySet_Contains(inlined_cells, name)) {
+ flags |= DEF_COMP_CELL;
+ }
  v_scope = PyDict_GetItemWithError(scopes, name);
  assert(v_scope && PyLong_Check(v_scope));
  scope = PyLong_AS_LONG(v_scope);
@@ -869,7 +870,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
  PySTEntryObject *class_entry)
 {
  PyObject *name, *v, *local = NULL, *scopes = NULL, *newbound = NULL;
- PyObject *newglobal = NULL, *newfree = NULL, *promote_to_cell = NULL;
+ PyObject *newglobal = NULL, *newfree = NULL, *inlined_cells = NULL;
  PyObject *temp;
  int success = 0;
  Py_ssize_t i, pos = 0;
@@ -901,8 +902,8 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
  newbound = PySet_New(NULL);
  if (!newbound)
  goto error;
- promote_to_cell = PySet_New(NULL);
- if (!promote_to_cell)
+ inlined_cells = PySet_New(NULL);
+ if (!inlined_cells)
  goto error;
 
  /* Class namespace has no effect on names visible in
@@ -996,7 +997,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
  goto error;
  }
  if (inline_comp) {
- if (!inline_comprehension(ste, entry, scopes, child_free, promote_to_cell)) {
+ if (!inline_comprehension(ste, entry, scopes, child_free, inlined_cells)) {
  Py_DECREF(child_free);
  goto error;
  }
@@ -1027,12 +1028,12 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
  }
 
  /* Check if any local variables must be converted to cell variables */
- if (_PyST_IsFunctionLike(ste) && !analyze_cells(scopes, newfree, promote_to_cell))
+ if (_PyST_IsFunctionLike(ste) && !analyze_cells(scopes, newfree, inlined_cells))
  goto error;
  else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree))
  goto error;
  /* Records the results of the analysis in the symbol table entry */
- if (!update_symbols(ste->ste_symbols, scopes, bound, newfree,
+ if (!update_symbols(ste->ste_symbols, scopes, bound, newfree, inlined_cells,
  ste->ste_type == ClassBlock))
  goto error;
 
@@ -1047,7 +1048,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
  Py_XDECREF(newbound);
  Py_XDECREF(newglobal);
  Py_XDECREF(newfree);
- Py_XDECREF(promote_to_cell);
+ Py_XDECREF(inlined_cells);
  if (!success)
  assert(PyErr_Occurred());
  return success;