gh-98442: fix locations of with statement's cleanup instructions (#120763)

gh-98442: fix location of with statement's cleanup instructions

1 files changed, 33 insertions, 0 deletions

diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
index a7c20cb..9def47e 100644
--- a/Lib/test/test_compile.py
+++ b/Lib/test/test_compile.py
@@ -2055,6 +2055,39 @@ class TestSourcePositions(unittest.TestCase):
                         self.assertGreaterEqual(end_col, start_col)
                         self.assertLessEqual(end_col, code_end)
 
+    def test_return_in_with_positions(self):
+        # See gh-98442
+        def f():
+            with xyz:
+                1
+                2
+                3
+                4
+                return R
+
+        # All instructions should have locations on a single line
+        for instr in dis.get_instructions(f):
+            start_line, end_line, _, _ = instr.positions
+            self.assertEqual(start_line, end_line)
+
+        # Expect three load None instructions for the no-exception __exit__ call,
+        # and one RETURN_VALUE.
+        # They should all have the locations of the context manager ('xyz').
+
+        load_none = [instr for instr in dis.get_instructions(f) if
+                     instr.opname == 'LOAD_CONST' and instr.argval is None]
+        return_value = [instr for instr in dis.get_instructions(f) if
+                        instr.opname == 'RETURN_VALUE']
+
+        self.assertEqual(len(load_none), 3)
+        self.assertEqual(len(return_value), 1)
+        for instr in load_none + return_value:
+            start_line, end_line, start_col, end_col = instr.positions
+            self.assertEqual(start_line, f.__code__.co_firstlineno + 1)
+            self.assertEqual(end_line, f.__code__.co_firstlineno + 1)
+            self.assertEqual(start_col, 17)
+            self.assertEqual(end_col, 20)
+
 
 class TestExpectedAttributes(unittest.TestCase):