Skip to content

Commit 86e6f16

Browse files
gh-104602: ensure all cellvars are known up front (#104603)
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
1 parent 3fadd7d commit 86e6f16

File tree

4 files changed

+50
-22
lines changed

4 files changed

+50
-22
lines changed

Include/internal/pycore_symtable.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -120,14 +120,15 @@ extern PyObject* _Py_Mangle(PyObject *p, PyObject *name);
120120
#define DEF_ANNOT 2<<7 /* this name is annotated */
121121
#define DEF_COMP_ITER 2<<8 /* this name is a comprehension iteration variable */
122122
#define DEF_TYPE_PARAM 2<<9 /* this name is a type parameter */
123+
#define DEF_COMP_CELL 2<<10 /* this name is a cell in an inlined comprehension */
123124

124125
#define DEF_BOUND (DEF_LOCAL | DEF_PARAM | DEF_IMPORT)
125126

126127
/* GLOBAL_EXPLICIT and GLOBAL_IMPLICIT are used internally by the symbol
127128
table. GLOBAL is returned from PyST_GetScope() for either of them.
128-
It is stored in ste_symbols at bits 12-15.
129+
It is stored in ste_symbols at bits 13-16.
129130
*/
130-
#define SCOPE_OFFSET 11
131+
#define SCOPE_OFFSET 12
131132
#define SCOPE_MASK (DEF_GLOBAL | DEF_LOCAL | DEF_PARAM | DEF_NONLOCAL)
132133

133134
#define LOCAL 1

Lib/test/test_listcomps.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,32 @@ def f():
381381
with self.assertRaises(UnboundLocalError):
382382
f()
383383

384+
def test_global_outside_cellvar_inside_plus_freevar(self):
385+
code = """
386+
a = 1
387+
def f():
388+
func, = [(lambda: b) for b in [a]]
389+
return b, func()
390+
x = f()
391+
"""
392+
self._check_in_scopes(
393+
code, {"x": (2, 1)}, ns={"b": 2}, scopes=["function", "module"])
394+
# inside a class, the `a = 1` assignment is not visible
395+
self._check_in_scopes(code, raises=NameError, scopes=["class"])
396+
397+
def test_cell_in_nested_comprehension(self):
398+
code = """
399+
a = 1
400+
def f():
401+
(func, inner_b), = [[lambda: b for b in c] + [b] for c in [[a]]]
402+
return b, inner_b, func()
403+
x = f()
404+
"""
405+
self._check_in_scopes(
406+
code, {"x": (2, 2, 1)}, ns={"b": 2}, scopes=["function", "module"])
407+
# inside a class, the `a = 1` assignment is not visible
408+
self._check_in_scopes(code, raises=NameError, scopes=["class"])
409+
384410
def test_name_error_in_class_scope(self):
385411
code = """
386412
y = 1

Python/compile.c

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1252,7 +1252,7 @@ compiler_enter_scope(struct compiler *c, identifier name,
12521252
}
12531253
u->u_metadata.u_name = Py_NewRef(name);
12541254
u->u_metadata.u_varnames = list2dict(u->u_ste->ste_varnames);
1255-
u->u_metadata.u_cellvars = dictbytype(u->u_ste->ste_symbols, CELL, 0, 0);
1255+
u->u_metadata.u_cellvars = dictbytype(u->u_ste->ste_symbols, CELL, DEF_COMP_CELL, 0);
12561256
if (!u->u_metadata.u_varnames || !u->u_metadata.u_cellvars) {
12571257
compiler_unit_free(u);
12581258
return ERROR;

Python/symtable.c

Lines changed: 20 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -632,7 +632,7 @@ is_free_in_any_child(PySTEntryObject *entry, PyObject *key)
632632
static int
633633
inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
634634
PyObject *scopes, PyObject *comp_free,
635-
PyObject *promote_to_cell)
635+
PyObject *inlined_cells)
636636
{
637637
PyObject *k, *v;
638638
Py_ssize_t pos = 0;
@@ -645,6 +645,11 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
645645
}
646646
int scope = (comp_flags >> SCOPE_OFFSET) & SCOPE_MASK;
647647
int only_flags = comp_flags & ((1 << SCOPE_OFFSET) - 1);
648+
if (scope == CELL || only_flags & DEF_COMP_CELL) {
649+
if (PySet_Add(inlined_cells, k) < 0) {
650+
return 0;
651+
}
652+
}
648653
PyObject *existing = PyDict_GetItemWithError(ste->ste_symbols, k);
649654
if (existing == NULL && PyErr_Occurred()) {
650655
return 0;
@@ -665,14 +670,6 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
665670
}
666671
else {
667672
if (PyLong_AsLong(existing) & DEF_BOUND) {
668-
// cell vars in comprehension that are locals in outer scope
669-
// must be promoted to cell so u_cellvars isn't wrong
670-
if (scope == CELL && _PyST_IsFunctionLike(ste)) {
671-
if (PySet_Add(promote_to_cell, k) < 0) {
672-
return 0;
673-
}
674-
}
675-
676673
// free vars in comprehension that are locals in outer scope can
677674
// now simply be locals, unless they are free in comp children,
678675
// or if the outer scope is a class block
@@ -698,7 +695,7 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
698695
*/
699696

700697
static int
701-
analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
698+
analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells)
702699
{
703700
PyObject *name, *v, *v_cell;
704701
int success = 0;
@@ -713,7 +710,7 @@ analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
713710
scope = PyLong_AS_LONG(v);
714711
if (scope != LOCAL)
715712
continue;
716-
if (!PySet_Contains(free, name) && !PySet_Contains(promote_to_cell, name))
713+
if (!PySet_Contains(free, name) && !PySet_Contains(inlined_cells, name))
717714
continue;
718715
/* Replace LOCAL with CELL for this name, and remove
719716
from free. It is safe to replace the value of name
@@ -753,7 +750,8 @@ drop_class_free(PySTEntryObject *ste, PyObject *free)
753750
*/
754751
static int
755752
update_symbols(PyObject *symbols, PyObject *scopes,
756-
PyObject *bound, PyObject *free, int classflag)
753+
PyObject *bound, PyObject *free,
754+
PyObject *inlined_cells, int classflag)
757755
{
758756
PyObject *name = NULL, *itr = NULL;
759757
PyObject *v = NULL, *v_scope = NULL, *v_new = NULL, *v_free = NULL;
@@ -764,6 +762,9 @@ update_symbols(PyObject *symbols, PyObject *scopes,
764762
long scope, flags;
765763
assert(PyLong_Check(v));
766764
flags = PyLong_AS_LONG(v);
765+
if (PySet_Contains(inlined_cells, name)) {
766+
flags |= DEF_COMP_CELL;
767+
}
767768
v_scope = PyDict_GetItemWithError(scopes, name);
768769
assert(v_scope && PyLong_Check(v_scope));
769770
scope = PyLong_AS_LONG(v_scope);
@@ -870,7 +871,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
870871
PySTEntryObject *class_entry)
871872
{
872873
PyObject *name, *v, *local = NULL, *scopes = NULL, *newbound = NULL;
873-
PyObject *newglobal = NULL, *newfree = NULL, *promote_to_cell = NULL;
874+
PyObject *newglobal = NULL, *newfree = NULL, *inlined_cells = NULL;
874875
PyObject *temp;
875876
int success = 0;
876877
Py_ssize_t i, pos = 0;
@@ -902,8 +903,8 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
902903
newbound = PySet_New(NULL);
903904
if (!newbound)
904905
goto error;
905-
promote_to_cell = PySet_New(NULL);
906-
if (!promote_to_cell)
906+
inlined_cells = PySet_New(NULL);
907+
if (!inlined_cells)
907908
goto error;
908909

909910
/* Class namespace has no effect on names visible in
@@ -997,7 +998,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
997998
goto error;
998999
}
9991000
if (inline_comp) {
1000-
if (!inline_comprehension(ste, entry, scopes, child_free, promote_to_cell)) {
1001+
if (!inline_comprehension(ste, entry, scopes, child_free, inlined_cells)) {
10011002
Py_DECREF(child_free);
10021003
goto error;
10031004
}
@@ -1028,12 +1029,12 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
10281029
}
10291030

10301031
/* Check if any local variables must be converted to cell variables */
1031-
if (_PyST_IsFunctionLike(ste) && !analyze_cells(scopes, newfree, promote_to_cell))
1032+
if (_PyST_IsFunctionLike(ste) && !analyze_cells(scopes, newfree, inlined_cells))
10321033
goto error;
10331034
else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree))
10341035
goto error;
10351036
/* Records the results of the analysis in the symbol table entry */
1036-
if (!update_symbols(ste->ste_symbols, scopes, bound, newfree,
1037+
if (!update_symbols(ste->ste_symbols, scopes, bound, newfree, inlined_cells,
10371038
ste->ste_type == ClassBlock))
10381039
goto error;
10391040

@@ -1048,7 +1049,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
10481049
Py_XDECREF(newbound);
10491050
Py_XDECREF(newglobal);
10501051
Py_XDECREF(newfree);
1051-
Py_XDECREF(promote_to_cell);
1052+
Py_XDECREF(inlined_cells);
10521053
if (!success)
10531054
assert(PyErr_Occurred());
10541055
return success;

0 commit comments

Comments
 (0)