Skip to content

Commit

Permalink
gh-104377: fix cell in comprehension that is free in outer scope (#10…
Browse files Browse the repository at this point in the history
  • Loading branch information
carljm committed May 11, 2023
1 parent 37a5d25 commit ac66cc1
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 7 deletions.
59 changes: 55 additions & 4 deletions Lib/test/test_listcomps.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,15 +117,15 @@ def get_output(moddict, name):
newcode = code
def get_output(moddict, name):
return moddict[name]
ns = ns or {}
newns = ns.copy() if ns else {}
try:
exec(newcode, ns)
exec(newcode, newns)
except raises as e:
# We care about e.g. NameError vs UnboundLocalError
self.assertIs(type(e), raises)
else:
for k, v in (outputs or {}).items():
self.assertEqual(get_output(ns, k), v)
self.assertEqual(get_output(newns, k), v)

def test_lambdas_with_iteration_var_as_default(self):
code = """
Expand Down Expand Up @@ -180,6 +180,26 @@ def test_closure_can_jump_over_comp_scope(self):
z = [x() for x in items]
"""
outputs = {"z": [2, 2, 2, 2, 2]}
self._check_in_scopes(code, outputs, scopes=["module", "function"])

def test_cell_inner_free_outer(self):
code = """
def f():
return [lambda: x for x in (x, [1])[1]]
x = ...
y = [fn() for fn in f()]
"""
outputs = {"y": [1]}
self._check_in_scopes(code, outputs, scopes=["module", "function"])

def test_free_inner_cell_outer(self):
code = """
g = 2
def f():
return g
y = [g for x in [1]]
"""
outputs = {"y": [2]}
self._check_in_scopes(code, outputs)

def test_inner_cell_shadows_outer_redefined(self):
Expand All @@ -203,6 +223,37 @@ def inner():
outputs = {"x": -1}
self._check_in_scopes(code, outputs, ns={"g": -1})

def test_explicit_global(self):
code = """
global g
x = g
g = 2
items = [g for g in [1]]
y = g
"""
outputs = {"x": 1, "y": 2, "items": [1]}
self._check_in_scopes(code, outputs, ns={"g": 1})

def test_explicit_global_2(self):
code = """
global g
x = g
g = 2
items = [g for x in [1]]
y = g
"""
outputs = {"x": 1, "y": 2, "items": [2]}
self._check_in_scopes(code, outputs, ns={"g": 1})

def test_explicit_global_3(self):
code = """
global g
fns = [lambda: g for g in [2]]
items = [fn() for fn in fns]
"""
outputs = {"items": [2]}
self._check_in_scopes(code, outputs, ns={"g": 1})

def test_assignment_expression(self):
code = """
x = -1
Expand Down Expand Up @@ -250,7 +301,7 @@ def g():
g()
"""
outputs = {"x": 1}
self._check_in_scopes(code, outputs)
self._check_in_scopes(code, outputs, scopes=["module", "function"])

def test_introspecting_frame_locals(self):
code = """
Expand Down
15 changes: 12 additions & 3 deletions Python/compile.c
Original file line number Diff line number Diff line change
Expand Up @@ -5028,14 +5028,19 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
long scope = (symbol >> SCOPE_OFFSET) & SCOPE_MASK;
PyObject *outv = PyDict_GetItemWithError(c->u->u_ste->ste_symbols, k);
if (outv == NULL) {
assert(PyErr_Occurred());
return ERROR;
}
assert(PyLong_Check(outv));
long outsc = (PyLong_AS_LONG(outv) >> SCOPE_OFFSET) & SCOPE_MASK;
if (scope != outsc) {
if (scope != outsc && !(scope == CELL && outsc == FREE)) {
// If a name has different scope inside than outside the
// comprehension, we need to temporarily handle it with the
// right scope while compiling the comprehension.
// right scope while compiling the comprehension. (If it's free
// in outer scope and cell in inner scope, we can't treat it as
// both cell and free in the same function, but treating it as
// free throughout is fine; it's *_DEREF either way.)

if (state->temp_symbols == NULL) {
state->temp_symbols = PyDict_New();
if (state->temp_symbols == NULL) {
Expand Down Expand Up @@ -5071,7 +5076,11 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
// comprehension and restore the original one after
ADDOP_NAME(c, loc, LOAD_FAST_AND_CLEAR, k, varnames);
if (scope == CELL) {
ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
if (outsc == FREE) {
ADDOP_NAME(c, loc, MAKE_CELL, k, freevars);
} else {
ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
}
}
if (PyList_Append(state->pushed_locals, k) < 0) {
return ERROR;
Expand Down

0 comments on commit ac66cc1

Please sign in to comment.