From 96c1737591b45bc1c528df1103eb0e500751fefe Mon Sep 17 00:00:00 2001
From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
Date: Thu, 22 Feb 2024 12:36:44 +0000
Subject: [PATCH] gh-115796: fix exception table construction in
 _testinternalcapi.assemble_code_object (#115797)

---
 Lib/test/test_compiler_assemble.py            | 41 ++++++++++++++++++-
 ...-02-22-00-17-06.gh-issue-115796.d4hpKy.rst |  2 +
 Python/flowgraph.c                            | 17 +++++---
 3 files changed, 52 insertions(+), 8 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Tests/2024-02-22-00-17-06.gh-issue-115796.d4hpKy.rst

diff --git a/Lib/test/test_compiler_assemble.py b/Lib/test/test_compiler_assemble.py
index 5696433e529..ab9f04dd63a 100644
--- a/Lib/test/test_compiler_assemble.py
+++ b/Lib/test/test_compiler_assemble.py
@@ -1,3 +1,6 @@
+import dis
+import io
+import textwrap
 import types
 
 from test.support.bytecode_helper import AssemblerTestCase
@@ -22,11 +25,13 @@ class IsolatedAssembleTests(AssemblerTestCase):
         metadata.setdefault('filename', filename)
         return metadata
 
-    def assemble_test(self, insts, metadata, expected):
+    def insts_to_code_object(self, insts, metadata):
         metadata = self.complete_metadata(metadata)
         insts = self.complete_insts_info(insts)
+        return self.get_code_object(metadata['filename'], insts, metadata)
 
-        co = self.get_code_object(metadata['filename'], insts, metadata)
+    def assemble_test(self, insts, metadata, expected):
+        co = self.insts_to_code_object(insts, metadata)
         self.assertIsInstance(co, types.CodeType)
 
         expected_metadata = {}
@@ -108,3 +113,35 @@ class IsolatedAssembleTests(AssemblerTestCase):
 
         expected = {(0,): 0, (1,): 1, (2,): 0, (120,): 0, (121,): 1}
         self.assemble_test(instructions, metadata, expected)
+
+
+    def test_exception_table(self):
+        metadata = {
+            'filename' : 'exc.py',
+            'name'     : 'exc',
+            'consts'   : {2 : 0},
+        }
+
+        # code for "try: pass\n except: pass"
+        insts = [
+            ('RESUME', 0),
+            ('SETUP_FINALLY', 3),
+            ('RETURN_CONST', 0),
+            ('SETUP_CLEANUP', 8),
+            ('PUSH_EXC_INFO', 0),
+            ('POP_TOP', 0),
+            ('POP_EXCEPT', 0),
+            ('RETURN_CONST', 0),
+            ('COPY', 3),
+            ('POP_EXCEPT', 0),
+            ('RERAISE', 1),
+        ]
+        co = self.insts_to_code_object(insts, metadata)
+        output = io.StringIO()
+        dis.dis(co, file=output)
+        exc_table = textwrap.dedent("""
+                                       ExceptionTable:
+                                         L1 to L2 -> L2 [0]
+                                         L2 to L3 -> L3 [1] lasti
+                                    """)
+        self.assertTrue(output.getvalue().endswith(exc_table))
diff --git a/Misc/NEWS.d/next/Tests/2024-02-22-00-17-06.gh-issue-115796.d4hpKy.rst b/Misc/NEWS.d/next/Tests/2024-02-22-00-17-06.gh-issue-115796.d4hpKy.rst
new file mode 100644
index 00000000000..a40be74f739
--- /dev/null
+++ b/Misc/NEWS.d/next/Tests/2024-02-22-00-17-06.gh-issue-115796.d4hpKy.rst
@@ -0,0 +1,2 @@
+Make '_testinternalcapi.assemble_code_object' construct the exception table
+for the code object.
diff --git a/Python/flowgraph.c b/Python/flowgraph.c
index 4d9ba9eceb8..2f47e47bf9d 100644
--- a/Python/flowgraph.c
+++ b/Python/flowgraph.c
@@ -665,12 +665,6 @@ translate_jump_labels_to_targets(basicblock *entryblock)
     return SUCCESS;
 }
 
-int
-_PyCfg_JumpLabelsToTargets(cfg_builder *g)
-{
-    return translate_jump_labels_to_targets(g->g_entryblock);
-}
-
 static int
 mark_except_handlers(basicblock *entryblock) {
 #ifndef NDEBUG
@@ -2790,3 +2784,14 @@ _PyCfg_OptimizedCfgToInstructionSequence(cfg_builder *g,
 
     return SUCCESS;
 }
+
+/* This is used by _PyCompile_Assemble to fill in the jump and exception
+ * targets in a synthetic CFG (which is not the ouptut of the builtin compiler).
+ */
+int
+_PyCfg_JumpLabelsToTargets(cfg_builder *g)
+{
+    RETURN_IF_ERROR(translate_jump_labels_to_targets(g->g_entryblock));
+    RETURN_IF_ERROR(label_exception_targets(g->g_entryblock));
+    return SUCCESS;
+}