Repository: python/cpython
Pull request: gh-117901: Add option for compiler's codegen to save nested instruction sequences for introspection
Fixed #117901.

<!-- gh-issue-number: gh-117901 -->
* Issue: gh-117901
<!-- /gh-issue-number -->


Problem statement:
_PyCompile_CodeGen compiles nested function all the way to code object
# Bug report

The ``_PyCompile_CodeGen`` test utility is supposed to run only codegen and return an instruction sequence. However, it compiles nested functions all the way to code object. Need to teach it how to return the instruction sequence for the nested function.

<!-- gh-linked-prs -->
### Linked PRs
* gh-118007
<!-- /gh-linked-prs -->

Hints from issue discussion:

Commit message: Merge branch 'main' into nested_functions

Patch to review:
diff --git a/Include/internal/pycore_instruction_sequence.h b/Include/internal/pycore_instruction_sequence.h
index ecba0d9d8e996e..d6a79616db71fa 100644
--- a/Include/internal/pycore_instruction_sequence.h
+++ b/Include/internal/pycore_instruction_sequence.h
@@ -61,6 +61,7 @@ _PyJumpTargetLabel _PyInstructionSequence_NewLabel(_PyInstructionSequence *seq);
 int _PyInstructionSequence_ApplyLabelMap(_PyInstructionSequence *seq);
 int _PyInstructionSequence_InsertInstruction(_PyInstructionSequence *seq, int pos,
                                              int opcode, int oparg, _Py_SourceLocation loc);
+int _PyInstructionSequence_AddNested(_PyInstructionSequence *seq, _PyInstructionSequence *nested);
 void PyInstructionSequence_Fini(_PyInstructionSequence *seq);
 
 extern PyTypeObject _PyInstructionSequence_Type;
diff --git a/Lib/test/test_compiler_codegen.py b/Lib/test/test_compiler_codegen.py
index 166294a40c1cb7..1088b4aa9e624d 100644
--- a/Lib/test/test_compiler_codegen.py
+++ b/Lib/test/test_compiler_codegen.py
@@ -1,4 +1,5 @@
 
+import textwrap
 from test.support.bytecode_helper import CodegenTestCase
 
 # Tests for the code-generation stage of the compiler.
@@ -6,11 +7,19 @@
 
 class IsolatedCodeGenTests(CodegenTestCase):
 
+    def assertInstructionsMatch_recursive(self, insts, expected_insts):
+        expected_nested = [i for i in expected_insts if isinstance(i, list)]
+        expected_insts = [i for i in expected_insts if not isinstance(i, list)]
+        self.assertInstructionsMatch(insts, expected_insts)
+        self.assertEqual(len(insts.get_nested()), len(expected_nested))
+        for n_insts, n_expected in zip(insts.get_nested(), expected_nested):
+            self.assertInstructionsMatch_recursive(n_insts, n_expected)
+
     def codegen_test(self, snippet, expected_insts):
         import ast
         a = ast.parse(snippet, "my_file.py", "exec")
         insts = self.generate_code(a)
-        self.assertInstructionsMatch(insts, expected_insts)
+        self.assertInstructionsMatch_recursive(insts, expected_insts)
 
     def test_if_expression(self):
         snippet = "42 if True else 24"
@@ -55,6 +64,91 @@ def test_for_loop(self):
         ]
         self.codegen_test(snippet, expected)
 
+    def test_function(self):
+        snippet = textwrap.dedent("""
+            def f(x):
+                return x + 42
+        """)
+        expected = [
+            # Function definition
+            ('RESUME', 0),
+            ('LOAD_CONST', 0),
+            ('MAKE_FUNCTION', None),
+            ('STORE_NAME', 0),
+            ('LOAD_CONST', 1),
+            ('RETURN_VALUE', None),
+            [
+                # Function body
+                ('RESUME', 0),
+                ('LOAD_FAST', 0),
+                ('LOAD_CONST', 1),
+                ('BINARY_OP', 0),
+                ('RETURN_VALUE', None),
+                ('LOAD_CONST', 0),
+                ('RETURN_VALUE', None),
+            ]
+        ]
+        self.codegen_test(snippet, expected)
+
+    def test_nested_functions(self):
+        snippet = textwrap.dedent("""
+            def f():
+                def h():
+                    return 12
+                def g():
+                    x = 1
+                    y = 2
+                    z = 3
+                    u = 4
+                    return 42
+        """)
+        expected = [
+            # Function definition
+            ('RESUME', 0),
+            ('LOAD_CONST', 0),
+            ('MAKE_FUNCTION', None),
+            ('STORE_NAME', 0),
+            ('LOAD_CONST', 1),
+            ('RETURN_VALUE', None),
+            [
+                # Function body
+                ('RESUME', 0),
+                ('LOAD_CONST', 1),
+                ('MAKE_FUNCTION', None),
+                ('STORE_FAST', 0),
+                ('LOAD_CONST', 2),
+                ('MAKE_FUNCTION', None),
+                ('STORE_FAST', 1),
+                ('LOAD_CONST', 0),
+                ('RETURN_VALUE', None),
+                [
+                    ('RESUME', 0),
+                    ('NOP', None),
+                    ('LOAD_CONST', 1),
+                    ('RETURN_VALUE', None),
+                    ('LOAD_CONST', 0),
+                    ('RETURN_VALUE', None),
+                ],
+                [
+                    ('RESUME', 0),
+                    ('LOAD_CONST', 1),
+                    ('STORE_FAST', 0),
+                    ('LOAD_CONST', 2),
+                    ('STORE_FAST', 1),
+                    ('LOAD_CONST', 3),
+                    ('STORE_FAST', 2),
+                    ('LOAD_CONST', 4),
+                    ('STORE_FAST', 3),
+                    ('NOP', None),
+                    ('LOAD_CONST', 5),
+                    ('RETURN_VALUE', None),
+                    ('LOAD_CONST', 0),
+                    ('RETURN_VALUE', None),
+                ],
+            ],
+        ]
+        self.codegen_test(snippet, expected)
+
     def test_syntax_error__return_not_in_function(self):
         snippet = "return 42"
         with self.assertRaisesRegex(SyntaxError, "'return' outside function"):
diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-04-17-22-53-52.gh-issue-117901.SsEcVJ.rst b/Misc/NEWS.d/next/Core and Builtins/2024-04-17-22-53-52.gh-issue-117901.SsEcVJ.rst
new file mode 100644
index 00000000000000..1e412690deecd7
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2024-04-17-22-53-52.gh-issue-117901.SsEcVJ.rst	
@@ -0,0 +1 @@
+Add option for compiler's codegen to save nested instruction sequences for introspection.
diff --git a/Python/compile.c b/Python/compile.c
index 3d856b7e4ddd97..b60931857426ce 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -285,6 +285,10 @@ struct compiler {
     struct compiler_unit *u;     /* compiler state for current block */
     PyObject *c_stack;           /* Python list holding compiler_unit ptrs */
     PyArena *c_arena;            /* pointer to memory allocation arena */
+
+    bool c_save_nested_seqs;     /* if true, construct recusrsive instruction sequences
+                                  * (including instructions for nested code objects)
+                                  */
 };
 
 #define INSTR_SEQUENCE(C) ((C)->u->u_instr_sequence)
@@ -402,6 +406,7 @@ compiler_setup(struct compiler *c, mod_ty mod, PyObject *filename,
     c->c_flags = *flags;
     c->c_optimize = (optimize == -1) ? _Py_GetConfig()->optimization_level : optimize;
     c->c_nestlevel = 0;
+    c->c_save_nested_seqs = false;
 
     if (!_PyAST_Optimize(mod, arena, c->c_optimize, merged)) {
         return ERROR;
@@ -1290,6 +1295,11 @@ compiler_exit_scope(struct compiler *c)
     // Don't call PySequence_DelItem() with an exception raised
     PyObject *exc = PyErr_GetRaisedException();
 
+    instr_sequence *nested_seq = NULL;
+    if (c->c_save_nested_seqs) {
+        nested_seq = c->u->u_instr_sequence;
+        Py_INCREF(nested_seq);
+    }
     c->c_nestlevel--;
     compiler_unit_free(c->u);
     /* Restore c->u to the parent unit. */
@@ -1303,10 +1313,17 @@ compiler_exit_scope(struct compiler *c)
             PyErr_FormatUnraisable("Exception ignored on removing "
                                    "the last compiler stack item");
         }
+        if (nested_seq != NULL) {
+            if (_PyInstructionSequence_AddNested(c->u->u_instr_sequence, nested_seq) < 0) {
+                PyErr_FormatUnraisable("Exception ignored on appending "
+                                       "nested instruction sequence");
+            }
+        }
     }
     else {
         c->u = NULL;
     }
+    Py_XDECREF(nested_seq);
 
     PyErr_SetRaisedException(exc);
 }
@@ -7734,6 +7751,7 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
         _PyArena_Free(arena);
         return NULL;
     }
+    c->c_save_nested_seqs = true;
 
     metadata = PyDict_New();
     if (metadata == NULL) {


Code review comments:
diff_hunk: @@ -285,6 +285,10 @@ struct compiler {
     struct compiler_unit *u;     /* compiler state for current block */
     PyObject *c_stack;           /* Python list holding compiler_unit ptrs */
     PyArena *c_arena;            /* pointer to memory allocation arena */
+
+    bool c_save_nested_seqs;     /* if true, construct recusrsive instruction sequences
path: Python/compile.c
text: @user1:
```suggestion
    bool c_save_nested_seqs;     /* if true, construct recursive instruction sequences
```