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

1 files changed, 159 insertions, 0 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.