gh-104404: fix crasher with nested comprehensions plus lambdas (#104442)

This commit is contained in:
Carl Meyer 2023-05-12 18:42:04 -06:00 committed by GitHub
parent 1eb950ca55
commit 563c7dcba0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 7 deletions

View file

@ -338,6 +338,14 @@ class ListComprehensionTest(unittest.TestCase):
outputs = {"y": [1, 3, 5]} outputs = {"y": [1, 3, 5]}
self._check_in_scopes(code, outputs) self._check_in_scopes(code, outputs)
def test_nested_4(self):
code = """
items = [([lambda: x for x in range(2)], lambda: x) for x in range(3)]
out = [([fn() for fn in fns], fn()) for fns, fn in items]
"""
outputs = {"out": [([1, 1], 2), ([1, 1], 2), ([1, 1], 2)]}
self._check_in_scopes(code, outputs)
def test_nameerror(self): def test_nameerror(self):
code = """ code = """
[x for x in [1]] [x for x in [1]]

View file

@ -575,7 +575,8 @@ is_free_in_any_child(PySTEntryObject *entry, PyObject *key)
static int static int
inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp, inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
PyObject *scopes, PyObject *comp_free) PyObject *scopes, PyObject *comp_free,
PyObject *promote_to_cell)
{ {
PyObject *k, *v; PyObject *k, *v;
Py_ssize_t pos = 0; Py_ssize_t pos = 0;
@ -611,7 +612,9 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
// cell vars in comprehension that are locals in outer scope // cell vars in comprehension that are locals in outer scope
// must be promoted to cell so u_cellvars isn't wrong // must be promoted to cell so u_cellvars isn't wrong
if (scope == CELL && ste->ste_type == FunctionBlock) { if (scope == CELL && ste->ste_type == FunctionBlock) {
SET_SCOPE(scopes, k, scope); if (PySet_Add(promote_to_cell, k) < 0) {
return 0;
}
} }
// free vars in comprehension that are locals in outer scope can // free vars in comprehension that are locals in outer scope can
@ -638,7 +641,7 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
*/ */
static int static int
analyze_cells(PyObject *scopes, PyObject *free) analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
{ {
PyObject *name, *v, *v_cell; PyObject *name, *v, *v_cell;
int success = 0; int success = 0;
@ -653,7 +656,7 @@ analyze_cells(PyObject *scopes, PyObject *free)
scope = PyLong_AS_LONG(v); scope = PyLong_AS_LONG(v);
if (scope != LOCAL) if (scope != LOCAL)
continue; continue;
if (!PySet_Contains(free, name)) if (!PySet_Contains(free, name) && !PySet_Contains(promote_to_cell, name))
continue; continue;
/* Replace LOCAL with CELL for this name, and remove /* Replace LOCAL with CELL for this name, and remove
from free. It is safe to replace the value of name from free. It is safe to replace the value of name
@ -803,7 +806,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
PyObject *global) PyObject *global)
{ {
PyObject *name, *v, *local = NULL, *scopes = NULL, *newbound = NULL; PyObject *name, *v, *local = NULL, *scopes = NULL, *newbound = NULL;
PyObject *newglobal = NULL, *newfree = NULL; PyObject *newglobal = NULL, *newfree = NULL, *promote_to_cell = NULL;
PyObject *temp; PyObject *temp;
int success = 0; int success = 0;
Py_ssize_t i, pos = 0; Py_ssize_t i, pos = 0;
@ -835,6 +838,9 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
newbound = PySet_New(NULL); newbound = PySet_New(NULL);
if (!newbound) if (!newbound)
goto error; goto error;
promote_to_cell = PySet_New(NULL);
if (!promote_to_cell)
goto error;
/* Class namespace has no effect on names visible in /* Class namespace has no effect on names visible in
nested functions, so populate the global and bound nested functions, so populate the global and bound
@ -915,7 +921,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
goto error; goto error;
} }
if (inline_comp) { if (inline_comp) {
if (!inline_comprehension(ste, entry, scopes, child_free)) { if (!inline_comprehension(ste, entry, scopes, child_free, promote_to_cell)) {
Py_DECREF(child_free); Py_DECREF(child_free);
goto error; goto error;
} }
@ -946,7 +952,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
} }
/* Check if any local variables must be converted to cell variables */ /* Check if any local variables must be converted to cell variables */
if (ste->ste_type == FunctionBlock && !analyze_cells(scopes, newfree)) if (ste->ste_type == FunctionBlock && !analyze_cells(scopes, newfree, promote_to_cell))
goto error; goto error;
else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree)) else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree))
goto error; goto error;
@ -966,6 +972,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
Py_XDECREF(newbound); Py_XDECREF(newbound);
Py_XDECREF(newglobal); Py_XDECREF(newglobal);
Py_XDECREF(newfree); Py_XDECREF(newfree);
Py_XDECREF(promote_to_cell);
if (!success) if (!success)
assert(PyErr_Occurred()); assert(PyErr_Occurred());
return success; return success;