Skip to content

Commit 800219c

Browse files
iritkatrieliritkatriel
committed
pythongh-98461: Fix source location in comprehensions bytecode
1 parent 1a6bacb commit 800219c

File tree

2 files changed

+108
-20
lines changed

2 files changed

+108
-20
lines changed

‎Lib/test/test_compile.py‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1249,6 +1249,86 @@ def test_multiline_assert(self):
12491249
self.assertOpcodeSourcePositionIs(compiled_code, 'RAISE_VARARGS',
12501250
line=1, end_line=3, column=0, end_column=30, occurrence=1)
12511251

1252+
def test_multiline_generator_expression(self):
1253+
snippet = """\
1254+
((x,
1255+
2*x)
1256+
for x
1257+
in [1,2,3] if (x > 0
1258+
and x < 100
1259+
and x != 50))
1260+
"""
1261+
1262+
compiled_code, _ = self.check_positions_against_ast(snippet)
1263+
compiled_code = compiled_code.co_consts[0]
1264+
self.assertIsInstance(compiled_code, types.CodeType)
1265+
self.assertOpcodeSourcePositionIs(compiled_code, 'YIELD_VALUE',
1266+
line=1, end_line=2, column=1, end_column=8, occurrence=1)
1267+
self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
1268+
line=1, end_line=2, column=1, end_column=8, occurrence=1)
1269+
self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
1270+
line=1, end_line=6, column=0, end_column=32, occurrence=1)
1271+
1272+
def test_multiline_list_comprehension(self):
1273+
snippet = """\
1274+
[(x,
1275+
2*x)
1276+
for x
1277+
in [1,2,3] if (x > 0
1278+
and x < 100
1279+
and x != 50)]
1280+
"""
1281+
1282+
compiled_code, _ = self.check_positions_against_ast(snippet)
1283+
compiled_code = compiled_code.co_consts[0]
1284+
self.assertIsInstance(compiled_code, types.CodeType)
1285+
self.assertOpcodeSourcePositionIs(compiled_code, 'LIST_APPEND',
1286+
line=1, end_line=2, column=1, end_column=8, occurrence=1)
1287+
self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
1288+
line=1, end_line=2, column=1, end_column=8, occurrence=1)
1289+
self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
1290+
line=1, end_line=6, column=0, end_column=32, occurrence=1)
1291+
1292+
def test_multiline_set_comprehension(self):
1293+
snippet = """\
1294+
{(x,
1295+
2*x)
1296+
for x
1297+
in [1,2,3] if (x > 0
1298+
and x < 100
1299+
and x != 50)}
1300+
"""
1301+
1302+
compiled_code, _ = self.check_positions_against_ast(snippet)
1303+
compiled_code = compiled_code.co_consts[0]
1304+
self.assertIsInstance(compiled_code, types.CodeType)
1305+
self.assertOpcodeSourcePositionIs(compiled_code, 'SET_ADD',
1306+
line=1, end_line=2, column=1, end_column=8, occurrence=1)
1307+
self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
1308+
line=1, end_line=2, column=1, end_column=8, occurrence=1)
1309+
self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
1310+
line=1, end_line=6, column=0, end_column=32, occurrence=1)
1311+
1312+
def test_multiline_dict_comprehension(self):
1313+
snippet = """\
1314+
{x:
1315+
2*x
1316+
for x
1317+
in [1,2,3] if (x > 0
1318+
and x < 100
1319+
and x != 50)}
1320+
"""
1321+
1322+
compiled_code, _ = self.check_positions_against_ast(snippet)
1323+
compiled_code = compiled_code.co_consts[0]
1324+
self.assertIsInstance(compiled_code, types.CodeType)
1325+
self.assertOpcodeSourcePositionIs(compiled_code, 'MAP_ADD',
1326+
line=1, end_line=2, column=1, end_column=7, occurrence=1)
1327+
self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
1328+
line=1, end_line=2, column=1, end_column=7, occurrence=1)
1329+
self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
1330+
line=1, end_line=6, column=0, end_column=32, occurrence=1)
1331+
12521332
def test_very_long_line_end_offset(self):
12531333
# Make sure we get the correct column offset for offsets
12541334
# too large to store in a byte.

‎Python/compile.c‎

Lines changed: 28 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -5208,21 +5208,19 @@ compiler_comprehension_generator(struct compiler *c, location *ploc,
52085208

52095209
static int
52105210
compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
5211-
asdl_comprehension_seq *generators, int gen_index,
5212-
int depth,
5211+
asdl_comprehension_seq *generators,
5212+
int gen_index, int depth,
52135213
expr_ty elt, expr_ty val, int type)
52145214
{
52155215
/* generate code for the iterator, then each of the ifs,
52165216
and then write to the element */
52175217

5218-
comprehension_ty gen;
5219-
Py_ssize_t i, n;
5220-
52215218
NEW_JUMP_TARGET_LABEL(c, start);
52225219
NEW_JUMP_TARGET_LABEL(c, if_cleanup);
52235220
NEW_JUMP_TARGET_LABEL(c, anchor);
52245221

5225-
gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
5222+
comprehension_ty gen = (comprehension_ty)asdl_seq_GET(generators,
5223+
gen_index);
52265224

52275225
if (gen_index == 0) {
52285226
/* Receive outermost iter as an implicit argument */
@@ -5265,42 +5263,51 @@ compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
52655263
VISIT(c, expr, gen->target);
52665264

52675265
/* XXX this needs to be cleaned up...a lot! */
5268-
n = asdl_seq_LEN(gen->ifs);
5269-
for (i = 0; i < n; i++) {
5266+
Py_ssize_t n = asdl_seq_LEN(gen->ifs);
5267+
for (Py_ssize_t i = 0; i < n; i++) {
52705268
expr_ty e = (expr_ty)asdl_seq_GET(gen->ifs, i);
5271-
if (!compiler_jump_if(c, ploc, e, if_cleanup, 0))
5269+
if (!compiler_jump_if(c, ploc, e, if_cleanup, 0)) {
52725270
return 0;
5271+
}
52735272
}
52745273

5275-
if (++gen_index < asdl_seq_LEN(generators))
5274+
if (++gen_index < asdl_seq_LEN(generators)) {
52765275
if (!compiler_comprehension_generator(c, ploc,
52775276
generators, gen_index, depth,
5278-
elt, val, type))
5279-
return 0;
5277+
elt, val, type)) {
5278+
return 0;
5279+
}
5280+
}
5281+
5282+
location elt_loc = LOC(elt);
52805283

52815284
/* only append after the last for generator */
52825285
if (gen_index >= asdl_seq_LEN(generators)) {
52835286
/* comprehension specific code */
52845287
switch (type) {
52855288
case COMP_GENEXP:
52865289
VISIT(c, expr, elt);
5287-
ADDOP_YIELD(c, *ploc);
5288-
ADDOP(c, *ploc, POP_TOP);
5290+
ADDOP_YIELD(c, elt_loc);
5291+
ADDOP(c, elt_loc, POP_TOP);
52895292
break;
52905293
case COMP_LISTCOMP:
52915294
VISIT(c, expr, elt);
5292-
ADDOP_I(c, *ploc, LIST_APPEND, depth + 1);
5295+
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
52935296
break;
52945297
case COMP_SETCOMP:
52955298
VISIT(c, expr, elt);
5296-
ADDOP_I(c, *ploc, SET_ADD, depth + 1);
5299+
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
52975300
break;
52985301
case COMP_DICTCOMP:
52995302
/* With '{k: v}', k is evaluated before v, so we do
53005303
the same. */
53015304
VISIT(c, expr, elt);
53025305
VISIT(c, expr, val);
5303-
ADDOP_I(c, *ploc, MAP_ADD, depth + 1);
5306+
elt_loc = LOCATION(elt->lineno,
5307+
val->end_lineno,
5308+
elt->col_offset,
5309+
val->end_col_offset);
5310+
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
53045311
break;
53055312
default:
53065313
return 0;
@@ -5309,7 +5316,7 @@ compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
53095316

53105317
USE_LABEL(c, if_cleanup);
53115318
if (IS_LABEL(start)) {
5312-
ADDOP_JUMP(c, *ploc, JUMP, start);
5319+
ADDOP_JUMP(c, elt_loc, JUMP, start);
53135320

53145321
USE_LABEL(c, anchor);
53155322
}
@@ -5467,11 +5474,12 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
54675474
}
54685475

54695476
if (!compiler_comprehension_generator(c, &loc, generators, 0, 0,
5470-
elt, val, type))
5477+
elt, val, type)) {
54715478
goto error_in_scope;
5479+
}
54725480

54735481
if (type != COMP_GENEXP) {
5474-
ADDOP(c, loc, RETURN_VALUE);
5482+
ADDOP(c, LOC(e), RETURN_VALUE);
54755483
}
54765484

54775485
co = assemble(c, 1);

0 commit comments

Comments
 (0)