gh-118513: Fix sibling comprehensions with a name bound in one and global in the other (#118526)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Co-authored-by: Kirill Podoprigora <kirill.bast9@mail.ru>
This commit is contained in:
Carl Meyer 2024-05-03 08:05:19 -06:00 committed by GitHub
parent 37ccf16786
commit c8deb1e4b4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 57 additions and 39 deletions

View File

@ -666,6 +666,20 @@ class ListComprehensionTest(unittest.TestCase):
self._check_in_scopes(code, expected)
self._check_in_scopes(code, expected, exec_func=self._replacing_exec)
def test_multiple_comprehension_name_reuse(self):
code = """
[x for x in [1]]
y = [x for _ in [1]]
"""
self._check_in_scopes(code, {"y": [3]}, ns={"x": 3})
code = """
x = 2
[x for x in [1]]
y = [x for _ in [1]]
"""
self._check_in_scopes(code, {"x": 2, "y": [3]}, ns={"x": 3}, scopes=["class"])
self._check_in_scopes(code, {"x": 2, "y": [2]}, ns={"x": 3}, scopes=["function", "module"])
__test__ = {'doctests' : doctests}

View File

@ -0,0 +1 @@
Fix incorrect :exc:`UnboundLocalError` when two comprehensions in the same function both reference the same name, and in one comprehension the name is bound while in the other it's an implicit global.

View File

@ -5561,10 +5561,48 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
while (PyDict_Next(entry->ste_symbols, &pos, &k, &v)) {
assert(PyLong_Check(v));
long symbol = PyLong_AS_LONG(v);
// only values bound in the comprehension (DEF_LOCAL) need to be handled
// at all; DEF_LOCAL | DEF_NONLOCAL can occur in the case of an
// assignment expression to a nonlocal in the comprehension, these don't
// need handling here since they shouldn't be isolated
long scope = (symbol >> SCOPE_OFFSET) & SCOPE_MASK;
PyObject *outv = PyDict_GetItemWithError(c->u->u_ste->ste_symbols, k);
if (outv == NULL) {
if (PyErr_Occurred()) {
return ERROR;
}
outv = _PyLong_GetZero();
}
assert(PyLong_CheckExact(outv));
long outsc = (PyLong_AS_LONG(outv) >> SCOPE_OFFSET) & SCOPE_MASK;
// 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. If it's free in the comprehension
// scope, no special handling; it should be handled the same as the
// enclosing scope. (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 ((scope != outsc && scope != FREE && !(scope == CELL && outsc == FREE))
|| in_class_block) {
if (state->temp_symbols == NULL) {
state->temp_symbols = PyDict_New();
if (state->temp_symbols == NULL) {
return ERROR;
}
}
// update the symbol to the in-comprehension version and save
// the outer version; we'll restore it after running the
// comprehension
Py_INCREF(outv);
if (PyDict_SetItem(c->u->u_ste->ste_symbols, k, v) < 0) {
Py_DECREF(outv);
return ERROR;
}
if (PyDict_SetItem(state->temp_symbols, k, outv) < 0) {
Py_DECREF(outv);
return ERROR;
}
Py_DECREF(outv);
}
// locals handling for names bound in comprehension (DEF_LOCAL |
// DEF_NONLOCAL occurs in assignment expression to nonlocal)
if ((symbol & DEF_LOCAL && !(symbol & DEF_NONLOCAL)) || in_class_block) {
if (!_PyST_IsFunctionLike(c->u->u_ste)) {
// non-function scope: override this name to use fast locals
@ -5589,41 +5627,6 @@ 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) {
outv = _PyLong_GetZero();
}
assert(PyLong_Check(outv));
long outsc = (PyLong_AS_LONG(outv) >> SCOPE_OFFSET) & SCOPE_MASK;
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. (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) {
return ERROR;
}
}
// update the symbol to the in-comprehension version and save
// the outer version; we'll restore it after running the
// comprehension
Py_INCREF(outv);
if (PyDict_SetItem(c->u->u_ste->ste_symbols, k, v) < 0) {
Py_DECREF(outv);
return ERROR;
}
if (PyDict_SetItem(state->temp_symbols, k, outv) < 0) {
Py_DECREF(outv);
return ERROR;
}
Py_DECREF(outv);
}
// local names bound in comprehension must be isolated from
// outer scope; push existing value (which may be NULL if
// not defined) on stack