gh-117901: Add option for compiler's codegen to save nested instruction sequences for introspection (#118007)

Author: iritkatriel

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;

Lib/test/test_compiler_codegen.py

@@ -1,16 +1,25 @@
 
+import textwrap
 from test.support.bytecode_helper import CodegenTestCase
 
 # Tests for the code-generation stage of the compiler.
 # Examine the un-optimized code generated from the AST.
 
 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"):

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.

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 recursive 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) {