From 80e0326b5380891638f75e0e0502039194336099 Mon Sep 17 00:00:00 2001
From: "Miss Islington (bot)"
 <31488909+miss-islington@users.noreply.github.com>
Date: Mon, 2 Sep 2024 17:46:41 +0200
Subject: [PATCH] [3.13] gh-98442: fix locations of with statement's cleanup
 instructions (GH-120763) (#120786)

gh-98442: fix locations of with statement's cleanup instructions (GH-120763)
(cherry picked from commit 55596ae0446e40f47e2a28b8897fe9530c32a19a)


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

Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
---
 Lib/test/test_compile.py                      | 33 +++++++++++++++++++
 ...4-06-19-21-34-21.gh-issue-98442.cqhjkN.rst |  2 ++
 Python/compile.c                              |  4 ++-
 3 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2024-06-19-21-34-21.gh-issue-98442.cqhjkN.rst

diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
index 4fa7e9aead0..b7e4d847822 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 TestStaticAttributes(unittest.TestCase):
 
diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-06-19-21-34-21.gh-issue-98442.cqhjkN.rst b/Misc/NEWS.d/next/Core and Builtins/2024-06-19-21-34-21.gh-issue-98442.cqhjkN.rst
new file mode 100644
index 00000000000..fb0a93f41a5
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2024-06-19-21-34-21.gh-issue-98442.cqhjkN.rst	
@@ -0,0 +1,2 @@
+Fix too wide source locations of the cleanup instructions of a with
+statement.
diff --git a/Python/compile.c b/Python/compile.c
index ed6923e1e5a..7752a6885cc 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -119,6 +119,7 @@ enum fblocktype { WHILE_LOOP, FOR_LOOP, TRY_EXCEPT, FINALLY_TRY, FINALLY_END,
 struct fblockinfo {
     enum fblocktype fb_type;
     jump_target_label fb_block;
+    location fb_loc;
     /* (optional) type-specific exit or cleanup block */
     jump_target_label fb_exit;
     /* (optional) additional information required for unwinding */
@@ -1433,6 +1434,7 @@ compiler_push_fblock(struct compiler *c, location loc,
     f = &c->u->u_fblock[c->u->u_nfblocks++];
     f->fb_type = t;
     f->fb_block = block_label;
+    f->fb_loc = loc;
     f->fb_exit = exit;
     f->fb_datum = datum;
     return SUCCESS;
@@ -1560,7 +1562,7 @@ compiler_unwind_fblock(struct compiler *c, location *ploc,
 
         case WITH:
         case ASYNC_WITH:
-            *ploc = LOC((stmt_ty)info->fb_datum);
+            *ploc = info->fb_loc;
             ADDOP(c, *ploc, POP_BLOCK);
             if (preserve_tos) {
                 ADDOP_I(c, *ploc, SWAP, 2);