gh-98461: Fix source location in comprehensions bytecode (GH-98464)

3 files changed, 234 insertions, 60 deletions

diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
index 85f05a8..a59d692 100644
--- a/Lib/test/test_compile.py
+++ b/Lib/test/test_compile.py
@@ -1249,6 +1249,165 @@ assert (a > 0 and
         self.assertOpcodeSourcePositionIs(compiled_code, 'RAISE_VARARGS',
             line=1, end_line=3, column=0, end_column=30, occurrence=1)
 
+    def test_multiline_generator_expression(self):
+        snippet = """\
+((x,
+    2*x)
+    for x
+    in [1,2,3] if (x > 0
+                   and x < 100
+                   and x != 50))
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        compiled_code = compiled_code.co_consts[0]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'YIELD_VALUE',
+            line=1, end_line=2, column=1, end_column=8, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
+            line=1, end_line=2, column=1, end_column=8, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=1, end_line=6, column=0, end_column=32, occurrence=1)
+
+    def test_multiline_async_generator_expression(self):
+        snippet = """\
+((x,
+    2*x)
+    async for x
+    in [1,2,3] if (x > 0
+                   and x < 100
+                   and x != 50))
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        compiled_code = compiled_code.co_consts[0]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'YIELD_VALUE',
+            line=1, end_line=2, column=1, end_column=8, occurrence=2)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=6, end_line=6, column=23, end_column=30, occurrence=1)
+
+    def test_multiline_list_comprehension(self):
+        snippet = """\
+[(x,
+    2*x)
+    for x
+    in [1,2,3] if (x > 0
+                   and x < 100
+                   and x != 50)]
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        compiled_code = compiled_code.co_consts[0]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'LIST_APPEND',
+            line=1, end_line=2, column=1, end_column=8, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
+            line=1, end_line=2, column=1, end_column=8, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=1, end_line=6, column=0, end_column=32, occurrence=1)
+
+    def test_multiline_async_list_comprehension(self):
+        snippet = """\
+async def f():
+    [(x,
+        2*x)
+        async for x
+        in [1,2,3] if (x > 0
+                       and x < 100
+                       and x != 50)]
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        g = {}
+        eval(compiled_code, g)
+        compiled_code = g['f'].__code__.co_consts[1]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'LIST_APPEND',
+            line=2, end_line=3, column=5, end_column=12, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
+            line=2, end_line=3, column=5, end_column=12, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=2, end_line=7, column=4, end_column=36, occurrence=1)
+
+    def test_multiline_set_comprehension(self):
+        snippet = """\
+{(x,
+    2*x)
+    for x
+    in [1,2,3] if (x > 0
+                   and x < 100
+                   and x != 50)}
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        compiled_code = compiled_code.co_consts[0]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'SET_ADD',
+            line=1, end_line=2, column=1, end_column=8, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
+            line=1, end_line=2, column=1, end_column=8, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=1, end_line=6, column=0, end_column=32, occurrence=1)
+
+    def test_multiline_async_set_comprehension(self):
+        snippet = """\
+async def f():
+    {(x,
+        2*x)
+        async for x
+        in [1,2,3] if (x > 0
+                       and x < 100
+                       and x != 50)}
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        g = {}
+        eval(compiled_code, g)
+        compiled_code = g['f'].__code__.co_consts[1]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'SET_ADD',
+            line=2, end_line=3, column=5, end_column=12, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
+            line=2, end_line=3, column=5, end_column=12, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=2, end_line=7, column=4, end_column=36, occurrence=1)
+
+    def test_multiline_dict_comprehension(self):
+        snippet = """\
+{x:
+    2*x
+    for x
+    in [1,2,3] if (x > 0
+                   and x < 100
+                   and x != 50)}
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        compiled_code = compiled_code.co_consts[0]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'MAP_ADD',
+            line=1, end_line=2, column=1, end_column=7, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
+            line=1, end_line=2, column=1, end_column=7, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=1, end_line=6, column=0, end_column=32, occurrence=1)
+
+    def test_multiline_async_dict_comprehension(self):
+        snippet = """\
+async def f():
+    {x:
+        2*x
+        async for x
+        in [1,2,3] if (x > 0
+                       and x < 100
+                       and x != 50)}
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        g = {}
+        eval(compiled_code, g)
+        compiled_code = g['f'].__code__.co_consts[1]
+        self.assertIsInstance(compiled_code, types.CodeType)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'MAP_ADD',
+            line=2, end_line=3, column=5, end_column=11, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
+            line=2, end_line=3, column=5, end_column=11, occurrence=1)
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=2, end_line=7, column=4, end_column=36, occurrence=1)
+
     def test_very_long_line_end_offset(self):
         # Make sure we get the correct column offset for offsets
         # too large to store in a byte.
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-10-19-20-53-38.gh-issue-98461.iNmPDV.rst b/Misc/NEWS.d/next/Core and Builtins/2022-10-19-20-53-38.gh-issue-98461.iNmPDV.rst
new file mode 100644
index 0000000..6289f20
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-10-19-20-53-38.gh-issue-98461.iNmPDV.rst
@@ -0,0 +1 @@
+Fix source location in bytecode for list, set and dict comprehensions as well as generator expressions.
diff --git a/Python/compile.c b/Python/compile.c
index c888f40..70f38d4 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -482,13 +482,13 @@ static int compiler_try_star_except(struct compiler *, stmt_ty);
 static int compiler_set_qualname(struct compiler *);
 
 static int compiler_sync_comprehension_generator(
-                                      struct compiler *c, location *ploc,
+                                      struct compiler *c, location loc,
                                       asdl_comprehension_seq *generators, int gen_index,
                                       int depth,
                                       expr_ty elt, expr_ty val, int type);
 
 static int compiler_async_comprehension_generator(
-                                      struct compiler *c, location *ploc,
+                                      struct compiler *c, location loc,
                                       asdl_comprehension_seq *generators, int gen_index,
                                       int depth,
                                       expr_ty elt, expr_ty val, int type);
@@ -5190,7 +5190,7 @@ ex_call:
 
 
 static int
-compiler_comprehension_generator(struct compiler *c, location *ploc,
+compiler_comprehension_generator(struct compiler *c, location loc,
                                  asdl_comprehension_seq *generators, int gen_index,
                                  int depth,
                                  expr_ty elt, expr_ty val, int type)
@@ -5199,35 +5199,33 @@ compiler_comprehension_generator(struct compiler *c, location *ploc,
     gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
     if (gen->is_async) {
         return compiler_async_comprehension_generator(
-            c, ploc, generators, gen_index, depth, elt, val, type);
+            c, loc, generators, gen_index, depth, elt, val, type);
     } else {
         return compiler_sync_comprehension_generator(
-            c, ploc, generators, gen_index, depth, elt, val, type);
+            c, loc, generators, gen_index, depth, elt, val, type);
     }
 }
 
 static int
-compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
-                                      asdl_comprehension_seq *generators, int gen_index,
-                                      int depth,
+compiler_sync_comprehension_generator(struct compiler *c, location loc,
+                                      asdl_comprehension_seq *generators,
+                                      int gen_index, int depth,
                                       expr_ty elt, expr_ty val, int type)
 {
     /* generate code for the iterator, then each of the ifs,
        and then write to the element */
 
-    comprehension_ty gen;
-    Py_ssize_t i, n;
-
     NEW_JUMP_TARGET_LABEL(c, start);
     NEW_JUMP_TARGET_LABEL(c, if_cleanup);
     NEW_JUMP_TARGET_LABEL(c, anchor);
 
-    gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
+    comprehension_ty gen = (comprehension_ty)asdl_seq_GET(generators,
+                                                          gen_index);
 
     if (gen_index == 0) {
         /* Receive outermost iter as an implicit argument */
         c->u->u_argcount = 1;
-        ADDOP_I(c, *ploc, LOAD_FAST, 0);
+        ADDOP_I(c, loc, LOAD_FAST, 0);
     }
     else {
         /* Sub-iter - calculate on the fly */
@@ -5254,29 +5252,34 @@ compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
         }
         if (IS_LABEL(start)) {
             VISIT(c, expr, gen->iter);
-            ADDOP(c, *ploc, GET_ITER);
+            ADDOP(c, loc, GET_ITER);
         }
     }
     if (IS_LABEL(start)) {
         depth++;
         USE_LABEL(c, start);
-        ADDOP_JUMP(c, *ploc, FOR_ITER, anchor);
+        ADDOP_JUMP(c, loc, FOR_ITER, anchor);
     }
     VISIT(c, expr, gen->target);
 
     /* XXX this needs to be cleaned up...a lot! */
-    n = asdl_seq_LEN(gen->ifs);
-    for (i = 0; i < n; i++) {
+    Py_ssize_t n = asdl_seq_LEN(gen->ifs);
+    for (Py_ssize_t i = 0; i < n; i++) {
         expr_ty e = (expr_ty)asdl_seq_GET(gen->ifs, i);
-        if (!compiler_jump_if(c, ploc, e, if_cleanup, 0))
+        if (!compiler_jump_if(c, &loc, e, if_cleanup, 0)) {
             return 0;
+        }
     }
 
-    if (++gen_index < asdl_seq_LEN(generators))
-        if (!compiler_comprehension_generator(c, ploc,
+    if (++gen_index < asdl_seq_LEN(generators)) {
+        if (!compiler_comprehension_generator(c, loc,
                                               generators, gen_index, depth,
-                                              elt, val, type))
-        return 0;
+                                              elt, val, type)) {
+            return 0;
+        }
+    }
+
+    location elt_loc = LOC(elt);
 
     /* only append after the last for generator */
     if (gen_index >= asdl_seq_LEN(generators)) {
@@ -5284,23 +5287,27 @@ compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
         switch (type) {
         case COMP_GENEXP:
             VISIT(c, expr, elt);
-            ADDOP_YIELD(c, *ploc);
-            ADDOP(c, *ploc, POP_TOP);
+            ADDOP_YIELD(c, elt_loc);
+            ADDOP(c, elt_loc, POP_TOP);
             break;
         case COMP_LISTCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, *ploc, LIST_APPEND, depth + 1);
+            ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
             break;
         case COMP_SETCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, *ploc, SET_ADD, depth + 1);
+            ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
             break;
         case COMP_DICTCOMP:
             /* With '{k: v}', k is evaluated before v, so we do
                the same. */
             VISIT(c, expr, elt);
             VISIT(c, expr, val);
-            ADDOP_I(c, *ploc, MAP_ADD, depth + 1);
+            elt_loc = LOCATION(elt->lineno,
+                               val->end_lineno,
+                               elt->col_offset,
+                               val->end_col_offset);
+            ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
             break;
         default:
             return 0;
@@ -5309,7 +5316,7 @@ compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
 
     USE_LABEL(c, if_cleanup);
     if (IS_LABEL(start)) {
-        ADDOP_JUMP(c, *ploc, JUMP, start);
+        ADDOP_JUMP(c, elt_loc, JUMP, start);
 
         USE_LABEL(c, anchor);
     }
@@ -5318,81 +5325,88 @@ compiler_sync_comprehension_generator(struct compiler *c, location *ploc,
 }
 
 static int
-compiler_async_comprehension_generator(struct compiler *c, location *ploc,
-                                      asdl_comprehension_seq *generators, int gen_index,
-                                      int depth,
+compiler_async_comprehension_generator(struct compiler *c, location loc,
+                                      asdl_comprehension_seq *generators,
+                                      int gen_index, int depth,
                                       expr_ty elt, expr_ty val, int type)
 {
-    comprehension_ty gen;
-    Py_ssize_t i, n;
     NEW_JUMP_TARGET_LABEL(c, start);
     NEW_JUMP_TARGET_LABEL(c, except);
     NEW_JUMP_TARGET_LABEL(c, if_cleanup);
 
-    gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
+    comprehension_ty gen = (comprehension_ty)asdl_seq_GET(generators,
+                                                          gen_index);
 
     if (gen_index == 0) {
         /* Receive outermost iter as an implicit argument */
         c->u->u_argcount = 1;
-        ADDOP_I(c, *ploc, LOAD_FAST, 0);
+        ADDOP_I(c, loc, LOAD_FAST, 0);
     }
     else {
         /* Sub-iter - calculate on the fly */
         VISIT(c, expr, gen->iter);
-        ADDOP(c, *ploc, GET_AITER);
+        ADDOP(c, loc, GET_AITER);
     }
 
     USE_LABEL(c, start);
     /* Runtime will push a block here, so we need to account for that */
-    if (!compiler_push_fblock(c, *ploc, ASYNC_COMPREHENSION_GENERATOR,
+    if (!compiler_push_fblock(c, loc, ASYNC_COMPREHENSION_GENERATOR,
                               start, NO_LABEL, NULL)) {
         return 0;
     }
 
-    ADDOP_JUMP(c, *ploc, SETUP_FINALLY, except);
-    ADDOP(c, *ploc, GET_ANEXT);
-    ADDOP_LOAD_CONST(c, *ploc, Py_None);
-    ADD_YIELD_FROM(c, *ploc, 1);
-    ADDOP(c, *ploc, POP_BLOCK);
+    ADDOP_JUMP(c, loc, SETUP_FINALLY, except);
+    ADDOP(c, loc, GET_ANEXT);
+    ADDOP_LOAD_CONST(c, loc, Py_None);
+    ADD_YIELD_FROM(c, loc, 1);
+    ADDOP(c, loc, POP_BLOCK);
     VISIT(c, expr, gen->target);
 
-    n = asdl_seq_LEN(gen->ifs);
-    for (i = 0; i < n; i++) {
+    Py_ssize_t n = asdl_seq_LEN(gen->ifs);
+    for (Py_ssize_t i = 0; i < n; i++) {
         expr_ty e = (expr_ty)asdl_seq_GET(gen->ifs, i);
-        if (!compiler_jump_if(c, ploc, e, if_cleanup, 0))
+        if (!compiler_jump_if(c, &loc, e, if_cleanup, 0)) {
             return 0;
+        }
     }
 
     depth++;
-    if (++gen_index < asdl_seq_LEN(generators))
-        if (!compiler_comprehension_generator(c, ploc,
+    if (++gen_index < asdl_seq_LEN(generators)) {
+        if (!compiler_comprehension_generator(c, loc,
                                               generators, gen_index, depth,
-                                              elt, val, type))
-        return 0;
+                                              elt, val, type)) {
+            return 0;
+        }
+    }
 
+    location elt_loc = LOC(elt);
     /* only append after the last for generator */
     if (gen_index >= asdl_seq_LEN(generators)) {
         /* comprehension specific code */
         switch (type) {
         case COMP_GENEXP:
             VISIT(c, expr, elt);
-            ADDOP_YIELD(c, *ploc);
-            ADDOP(c, *ploc, POP_TOP);
+            ADDOP_YIELD(c, elt_loc);
+            ADDOP(c, elt_loc, POP_TOP);
             break;
         case COMP_LISTCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, *ploc, LIST_APPEND, depth + 1);
+            ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
             break;
         case COMP_SETCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, *ploc, SET_ADD, depth + 1);
+            ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
             break;
         case COMP_DICTCOMP:
             /* With '{k: v}', k is evaluated before v, so we do
                the same. */
             VISIT(c, expr, elt);
             VISIT(c, expr, val);
-            ADDOP_I(c, *ploc, MAP_ADD, depth + 1);
+            elt_loc = LOCATION(elt->lineno,
+                               val->end_lineno,
+                               elt->col_offset,
+                               val->end_col_offset);
+            ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
             break;
         default:
             return 0;
@@ -5400,14 +5414,13 @@ compiler_async_comprehension_generator(struct compiler *c, location *ploc,
     }
 
     USE_LABEL(c, if_cleanup);
-    ADDOP_JUMP(c, *ploc, JUMP, start);
+    ADDOP_JUMP(c, elt_loc, JUMP, start);
 
     compiler_pop_fblock(c, ASYNC_COMPREHENSION_GENERATOR, start);
 
     USE_LABEL(c, except);
-    //UNSET_LOC(c);
 
-    ADDOP(c, *ploc, END_ASYNC_FOR);
+    ADDOP(c, loc, END_ASYNC_FOR);
 
     return 1;
 }
@@ -5466,12 +5479,13 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
         ADDOP_I(c, loc, op, 0);
     }
 
-    if (!compiler_comprehension_generator(c, &loc, generators, 0, 0,
-                                          elt, val, type))
+    if (!compiler_comprehension_generator(c, loc, generators, 0, 0,
+                                          elt, val, type)) {
         goto error_in_scope;
+    }
 
     if (type != COMP_GENEXP) {
-        ADDOP(c, loc, RETURN_VALUE);
+        ADDOP(c, LOC(e), RETURN_VALUE);
     }
 
     co = assemble(c, 1);