gh-93691: fix too broad source locations of with-statement instructions (#120125)

3 files changed, 47 insertions, 3 deletions

diff --git a/Lib/test/test_with.py b/Lib/test/test_with.py
index d819023..e8c4ddf 100644
--- a/Lib/test/test_with.py
+++ b/Lib/test/test_with.py
@@ -5,6 +5,7 @@ __author__ = "Mike Bland"
 __email__ = "mbland at acm dot org"
 
 import sys
+import traceback
 import unittest
 from collections import deque
 from contextlib import _GeneratorContextManager, contextmanager, nullcontext
@@ -749,5 +750,48 @@ class NestedWith(unittest.TestCase):
             self.assertEqual(10, b1)
             self.assertEqual(20, b2)
 
+    def testExceptionLocation(self):
+        # The location of an exception raised from
+        # __init__, __enter__ or __exit__ of a context
+        # manager should be just the context manager expression,
+        # pinpointing the precise context manager in case there
+        # is more than one.
+
+        def init_raises():
+            try:
+                with self.Dummy(), self.InitRaises() as cm, self.Dummy() as d:
+                    pass
+            except Exception as e:
+                return e
+
+        def enter_raises():
+            try:
+                with self.EnterRaises(), self.Dummy() as d:
+                    pass
+            except Exception as e:
+                return e
+
+        def exit_raises():
+            try:
+                with self.ExitRaises(), self.Dummy() as d:
+                    pass
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "self.InitRaises()"),
+                               (enter_raises, "self.EnterRaises()"),
+                               (exit_raises, "self.ExitRaises()"),
+                              ]:
+            with self.subTest(func):
+                exc = func()
+                f = traceback.extract_tb(exc.__traceback__)[0]
+                indent = 16
+                co = func.__code__
+                self.assertEqual(f.lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.end_lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.line[f.colno - indent : f.end_colno - indent],
+                                 expected)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-06-05-18-29-18.gh-issue-93691.6OautB.rst b/Misc/NEWS.d/next/Core and Builtins/2024-06-05-18-29-18.gh-issue-93691.6OautB.rst
new file mode 100644
index 0000000..c06d5a2
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2024-06-05-18-29-18.gh-issue-93691.6OautB.rst
@@ -0,0 +1 @@
+Fix source locations of instructions generated for with statements.
diff --git a/Python/compile.c b/Python/compile.c
index 7d74096..cb72415 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -5900,7 +5900,7 @@ compiler_async_with(struct compiler *c, stmt_ty s, int pos)
 
     /* Evaluate EXPR */
     VISIT(c, expr, item->context_expr);
-
+    loc = LOC(item->context_expr);
     ADDOP(c, loc, BEFORE_ASYNC_WITH);
     ADDOP_I(c, loc, GET_AWAITABLE, 1);
     ADDOP_LOAD_CONST(c, loc, Py_None);
@@ -5998,7 +5998,7 @@ compiler_with(struct compiler *c, stmt_ty s, int pos)
     /* Evaluate EXPR */
     VISIT(c, expr, item->context_expr);
     /* Will push bound __exit__ */
-    location loc = LOC(s);
+    location loc = LOC(item->context_expr);
     ADDOP(c, loc, BEFORE_WITH);
     ADDOP_JUMP(c, loc, SETUP_WITH, final);
 
@@ -6031,7 +6031,6 @@ compiler_with(struct compiler *c, stmt_ty s, int pos)
     /* For successful outcome:
      * call __exit__(None, None, None)
      */
-    loc = LOC(s);
     RETURN_IF_ERROR(compiler_call_exit_with_nones(c, loc));
     ADDOP(c, loc, POP_TOP);
     ADDOP_JUMP(c, loc, JUMP, exit);