From 0aa0fc3d3ca144f979c684552a56a18ed8f558e4 Mon Sep 17 00:00:00 2001
From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
Date: Wed, 24 Apr 2024 10:46:17 +0100
Subject: gh-117901: Add option for compiler's codegen to save nested
 instruction sequences for introspection (#118007)

---
 Lib/test/test_compiler_codegen.py | 96 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 95 insertions(+), 1 deletion(-)

(limited to 'Lib/test/test_compiler_codegen.py')

diff --git a/Lib/test/test_compiler_codegen.py b/Lib/test/test_compiler_codegen.py
index 166294a40c1..1088b4aa9e6 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 @@ from test.support.bytecode_helper import CodegenTestCase
 
 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 @@ class IsolatedCodeGenTests(CodegenTestCase):
         ]
         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"):
-- 
cgit v1.2.3