Skip to content

Commit 662aede

Browse files
gh-104374: Remove access to class scopes for inlined comprehensions (#104528)
Co-authored-by: Carl Meyer <carl@oddbird.net>
1 parent 152227b commit 662aede

File tree

3 files changed

+125
-9
lines changed

3 files changed

+125
-9
lines changed

Lib/test/test_listcomps.py

+107-2
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,8 @@ def f():
200200
y = [g for x in [1]]
201201
"""
202202
outputs = {"y": [2]}
203-
self._check_in_scopes(code, outputs)
203+
self._check_in_scopes(code, outputs, scopes=["module", "function"])
204+
self._check_in_scopes(code, scopes=["class"], raises=NameError)
204205

205206
def test_inner_cell_shadows_outer_redefined(self):
206207
code = """
@@ -328,7 +329,8 @@ def test_nested_2(self):
328329
y = [x for [x ** x for x in range(x)][x - 1] in l]
329330
"""
330331
outputs = {"y": [3, 3, 3]}
331-
self._check_in_scopes(code, outputs)
332+
self._check_in_scopes(code, outputs, scopes=["module", "function"])
333+
self._check_in_scopes(code, scopes=["class"], raises=NameError)
332334

333335
def test_nested_3(self):
334336
code = """
@@ -379,6 +381,109 @@ def f():
379381
with self.assertRaises(UnboundLocalError):
380382
f()
381383

384+
def test_name_error_in_class_scope(self):
385+
code = """
386+
y = 1
387+
[x + y for x in range(2)]
388+
"""
389+
self._check_in_scopes(code, raises=NameError, scopes=["class"])
390+
391+
def test_global_in_class_scope(self):
392+
code = """
393+
y = 2
394+
vals = [(x, y) for x in range(2)]
395+
"""
396+
outputs = {"vals": [(0, 1), (1, 1)]}
397+
self._check_in_scopes(code, outputs, ns={"y": 1}, scopes=["class"])
398+
399+
def test_in_class_scope_inside_function_1(self):
400+
code = """
401+
class C:
402+
y = 2
403+
vals = [(x, y) for x in range(2)]
404+
vals = C.vals
405+
"""
406+
outputs = {"vals": [(0, 1), (1, 1)]}
407+
self._check_in_scopes(code, outputs, ns={"y": 1}, scopes=["function"])
408+
409+
def test_in_class_scope_inside_function_2(self):
410+
code = """
411+
y = 1
412+
class C:
413+
y = 2
414+
vals = [(x, y) for x in range(2)]
415+
vals = C.vals
416+
"""
417+
outputs = {"vals": [(0, 1), (1, 1)]}
418+
self._check_in_scopes(code, outputs, scopes=["function"])
419+
420+
def test_in_class_scope_with_global(self):
421+
code = """
422+
y = 1
423+
class C:
424+
global y
425+
y = 2
426+
# Ensure the listcomp uses the global, not the value in the
427+
# class namespace
428+
locals()['y'] = 3
429+
vals = [(x, y) for x in range(2)]
430+
vals = C.vals
431+
"""
432+
outputs = {"vals": [(0, 2), (1, 2)]}
433+
self._check_in_scopes(code, outputs, scopes=["module", "class"])
434+
outputs = {"vals": [(0, 1), (1, 1)]}
435+
self._check_in_scopes(code, outputs, scopes=["function"])
436+
437+
def test_in_class_scope_with_nonlocal(self):
438+
code = """
439+
y = 1
440+
class C:
441+
nonlocal y
442+
y = 2
443+
# Ensure the listcomp uses the global, not the value in the
444+
# class namespace
445+
locals()['y'] = 3
446+
vals = [(x, y) for x in range(2)]
447+
vals = C.vals
448+
"""
449+
outputs = {"vals": [(0, 2), (1, 2)]}
450+
self._check_in_scopes(code, outputs, scopes=["function"])
451+
452+
def test_nested_has_free_var(self):
453+
code = """
454+
items = [a for a in [1] if [a for _ in [0]]]
455+
"""
456+
outputs = {"items": [1]}
457+
self._check_in_scopes(code, outputs, scopes=["class"])
458+
459+
def test_nested_free_var_not_bound_in_outer_comp(self):
460+
code = """
461+
z = 1
462+
items = [a for a in [1] if [x for x in [1] if z]]
463+
"""
464+
self._check_in_scopes(code, {"items": [1]}, scopes=["module", "function"])
465+
self._check_in_scopes(code, {"items": []}, ns={"z": 0}, scopes=["class"])
466+
467+
def test_nested_free_var_in_iter(self):
468+
code = """
469+
items = [_C for _C in [1] for [0, 1][[x for x in [1] if _C][0]] in [2]]
470+
"""
471+
self._check_in_scopes(code, {"items": [1]})
472+
473+
def test_nested_free_var_in_expr(self):
474+
code = """
475+
items = [(_C, [x for x in [1] if _C]) for _C in [0, 1]]
476+
"""
477+
self._check_in_scopes(code, {"items": [(0, []), (1, [1])]})
478+
479+
def test_nested_listcomp_in_lambda(self):
480+
code = """
481+
f = [(z, lambda y: [(x, y, z) for x in [3]]) for z in [1]]
482+
(z, func), = f
483+
out = func(2)
484+
"""
485+
self._check_in_scopes(code, {"z": 1, "out": [(3, 2, 1)]})
486+
382487

383488
__test__ = {'doctests' : doctests}
384489

Python/compile.c

+15-5
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,8 @@ struct compiler_unit {
388388
instr_sequence u_instr_sequence; /* codegen output */
389389

390390
int u_nfblocks;
391+
int u_in_inlined_comp;
392+
391393
struct fblockinfo u_fblock[CO_MAXBLOCKS];
392394

393395
_PyCompile_CodeUnitMetadata u_metadata;
@@ -1290,6 +1292,7 @@ compiler_enter_scope(struct compiler *c, identifier name,
12901292
}
12911293

12921294
u->u_nfblocks = 0;
1295+
u->u_in_inlined_comp = 0;
12931296
u->u_metadata.u_firstlineno = lineno;
12941297
u->u_metadata.u_consts = PyDict_New();
12951298
if (!u->u_metadata.u_consts) {
@@ -4137,7 +4140,7 @@ compiler_nameop(struct compiler *c, location loc,
41374140
case OP_DEREF:
41384141
switch (ctx) {
41394142
case Load:
4140-
if (c->u->u_ste->ste_type == ClassBlock) {
4143+
if (c->u->u_ste->ste_type == ClassBlock && !c->u->u_in_inlined_comp) {
41414144
op = LOAD_FROM_DICT_OR_DEREF;
41424145
// First load the locals
41434146
if (codegen_addop_noarg(INSTR_SEQUENCE(c), LOAD_LOCALS, loc) < 0) {
@@ -4188,7 +4191,12 @@ compiler_nameop(struct compiler *c, location loc,
41884191
break;
41894192
case OP_NAME:
41904193
switch (ctx) {
4191-
case Load: op = LOAD_NAME; break;
4194+
case Load:
4195+
op = (c->u->u_ste->ste_type == ClassBlock
4196+
&& c->u->u_in_inlined_comp)
4197+
? LOAD_GLOBAL
4198+
: LOAD_NAME;
4199+
break;
41924200
case Store: op = STORE_NAME; break;
41934201
case Del: op = DELETE_NAME; break;
41944202
}
@@ -5415,6 +5423,8 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
54155423
PySTEntryObject *entry,
54165424
inlined_comprehension_state *state)
54175425
{
5426+
int in_class_block = (c->u->u_ste->ste_type == ClassBlock) && !c->u->u_in_inlined_comp;
5427+
c->u->u_in_inlined_comp++;
54185428
// iterate over names bound in the comprehension and ensure we isolate
54195429
// them from the outer scope as needed
54205430
PyObject *k, *v;
@@ -5426,7 +5436,7 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
54265436
// at all; DEF_LOCAL | DEF_NONLOCAL can occur in the case of an
54275437
// assignment expression to a nonlocal in the comprehension, these don't
54285438
// need handling here since they shouldn't be isolated
5429-
if (symbol & DEF_LOCAL && !(symbol & DEF_NONLOCAL)) {
5439+
if ((symbol & DEF_LOCAL && !(symbol & DEF_NONLOCAL)) || in_class_block) {
54305440
if (!_PyST_IsFunctionLike(c->u->u_ste)) {
54315441
// non-function scope: override this name to use fast locals
54325442
PyObject *orig = PyDict_GetItem(c->u->u_metadata.u_fasthidden, k);
@@ -5448,8 +5458,7 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
54485458
long scope = (symbol >> SCOPE_OFFSET) & SCOPE_MASK;
54495459
PyObject *outv = PyDict_GetItemWithError(c->u->u_ste->ste_symbols, k);
54505460
if (outv == NULL) {
5451-
assert(PyErr_Occurred());
5452-
return ERROR;
5461+
outv = _PyLong_GetZero();
54535462
}
54545463
assert(PyLong_Check(outv));
54555464
long outsc = (PyLong_AS_LONG(outv) >> SCOPE_OFFSET) & SCOPE_MASK;
@@ -5523,6 +5532,7 @@ static int
55235532
pop_inlined_comprehension_state(struct compiler *c, location loc,
55245533
inlined_comprehension_state state)
55255534
{
5535+
c->u->u_in_inlined_comp--;
55265536
PyObject *k, *v;
55275537
Py_ssize_t pos = 0;
55285538
if (state.temp_symbols) {

Python/symtable.c

+3-2
Original file line numberDiff line numberDiff line change
@@ -674,8 +674,9 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
674674
}
675675

676676
// free vars in comprehension that are locals in outer scope can
677-
// now simply be locals, unless they are free in comp children
678-
if (!is_free_in_any_child(comp, k)) {
677+
// now simply be locals, unless they are free in comp children,
678+
// or if the outer scope is a class block
679+
if (!is_free_in_any_child(comp, k) && ste->ste_type != ClassBlock) {
679680
if (PySet_Discard(comp_free, k) < 0) {
680681
return 0;
681682
}

0 commit comments

Comments
 (0)