gh-104357: fix inlined comprehensions that close over iteration var (#104368)

This commit is contained in:
Carl Meyer 2023-05-10 19:08:40 -06:00 committed by GitHub
parent 94f30c7557
commit fcd5fb49b1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 23 additions and 6 deletions

View file

@ -163,6 +163,16 @@ class ListComprehensionTest(unittest.TestCase):
outputs = {"y": [4, 4, 4, 4, 4], "i": 20}
self._check_in_scopes(code, outputs)
def test_inner_cell_shadows_outer_no_store(self):
code = """
def f(x):
return [lambda: x for x in range(x)], x
fns, x = f(2)
y = [fn() for fn in fns]
"""
outputs = {"y": [1, 1], "x": 2}
self._check_in_scopes(code, outputs)
def test_closure_can_jump_over_comp_scope(self):
code = """
items = [(lambda: y) for i in range(5)]

View file

@ -607,12 +607,19 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
SET_SCOPE(scopes, k, scope);
}
else {
// free vars in comprehension that are locals in outer scope can
// now simply be locals, unless they are free in comp children
if ((PyLong_AsLong(existing) & DEF_BOUND) &&
!is_free_in_any_child(comp, k)) {
if (PySet_Discard(comp_free, k) < 0) {
return 0;
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 && ste->ste_type == FunctionBlock) {
SET_SCOPE(scopes, k, scope);
}
// 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)) {
if (PySet_Discard(comp_free, k) < 0) {
return 0;
}
}
}
}