gh-122666: Tests for ast optimizations

Lib/test/test_ast/snippets.py

@@ -366,23 +366,72 @@
   "f'foo({a})'",
 ]
 
+eval_opt_tests = [
+    # Common arithmetic operations
+    # There's no snippet for `@` operand since we don't have
+    # builtin constant that implements `@` operator
+    "1 + 1",
+    "1 - 1",
+    "1 * 1",
+    "1 / 1",
+    "1 // 1",
+    "1 % 1",
+    "1 ** 1",
+    "1 << 1",
+    "1 >> 1",
+    "1 | 1",
+    "1 ^ 1",
+    "1 & 1",
+    # Unary operators
+    "~1",
+    "-1",
+    "+1",
+    # Transformation from list to tuple and from set to frozenset
+    "1 in [1, 2, 3]",
+    "1 in {1, 2, 3}",
+    # Tuple must be folded
+    "(1, 2, 3)",
+    # Subscription
+    "(1, 2, 3)[0]",
+]
+
+exec_opt_tests = [
+    # The same transformation as in eval_opt_tests
+    "for _ in [1, 2, 3]:\n pass",
+    "for _ in {1, 2, 3}:\n pass",
+]
+
 
 def main():
     if __name__ != '__main__':
         return
-    if sys.argv[1:] == ['-g']:
+    args = sys.argv[1:]
+    if args == ["-g"]:
         for statements, kind in ((exec_tests, "exec"), (single_tests, "single"),
                                  (eval_tests, "eval")):
             print(kind+"_results = [")
             for statement in statements:
                 tree = ast.parse(statement, "?", kind)
                 print("%r," % (to_tuple(tree),))
             print("]")
-        print("main()")
-        raise SystemExit
-    unittest.main()
+    elif args == ["-g_opt"]:
+        for snippets, kind in ((eval_opt_tests, "eval"), (exec_opt_tests, "exec")):
+            folded = kind+"_results_folded = ["
+            not_folded = kind+"_results_not_folded =["
+            for snippet in snippets:
+                folded_tree = ast.parse(snippet, "?", kind, optimize=2)
+                not_folded_tree = ast.parse(snippet, "?", kind, optimize=-1)
+
+                folded += "%r," % (to_tuple(folded_tree),) + "\n"
+                not_folded += "%r," % (to_tuple(not_folded_tree),) + "\n"
+            folded += "]"
+            not_folded += "]"
+            print(folded)
+            print(not_folded)
+
+
 
-#### EVERYTHING BELOW IS GENERATED BY python Lib/test/test_ast/snippets.py -g  #####
+# Lines below are generated by running Lib/test/test_ast/snippets.py -g  #####
 exec_results = [
 ('Module', [('Expr', (1, 0, 1, 18), ('Constant', (1, 0, 1, 18), 'module docstring', None))], []),
 ('Module', [('FunctionDef', (1, 0, 1, 13), 'f', ('arguments', [], [], None, [], [], None, []), [('Pass', (1, 9, 1, 13))], [], None, None, [])], []),
@@ -599,4 +648,53 @@ def main():
 ('Expression', ('JoinedStr', (1, 0, 1, 8), [('FormattedValue', (1, 2, 1, 7), ('Name', (1, 3, 1, 4), 'a', ('Load',)), 114, None)])),
 ('Expression', ('JoinedStr', (1, 0, 1, 11), [('Constant', (1, 2, 1, 6), 'foo(', None), ('FormattedValue', (1, 6, 1, 9), ('Name', (1, 7, 1, 8), 'a', ('Load',)), -1, None), ('Constant', (1, 9, 1, 10), ')', None)])),
 ]
+
+# Lines below are generated by running Lib/test/test_ast/snippets.py -g_opt
+eval_results_folded = [('Expression', ('Constant', (1, 0, 1, 5), 2, None)),
+('Expression', ('Constant', (1, 0, 1, 5), 0, None)),
+('Expression', ('Constant', (1, 0, 1, 5), 1, None)),
+('Expression', ('Constant', (1, 0, 1, 5), 1.0, None)),
+('Expression', ('Constant', (1, 0, 1, 6), 1, None)),
+('Expression', ('Constant', (1, 0, 1, 5), 0, None)),
+('Expression', ('Constant', (1, 0, 1, 6), 1, None)),
+('Expression', ('Constant', (1, 0, 1, 6), 2, None)),
+('Expression', ('Constant', (1, 0, 1, 6), 0, None)),
+('Expression', ('Constant', (1, 0, 1, 5), 1, None)),
+('Expression', ('Constant', (1, 0, 1, 5), 0, None)),
+('Expression', ('Constant', (1, 0, 1, 5), 1, None)),
+('Expression', ('Constant', (1, 0, 1, 2), -2, None)),
+('Expression', ('Constant', (1, 0, 1, 2), -1, None)),
+('Expression', ('Constant', (1, 0, 1, 2), 1, None)),
+('Expression', ('Compare', (1, 0, 1, 14), ('Constant', (1, 0, 1, 1), 1, None), [('In',)], [('Constant', (1, 5, 1, 14), (1, 2, 3), None)])),
+('Expression', ('Compare', (1, 0, 1, 14), ('Constant', (1, 0, 1, 1), 1, None), [('In',)], [('Constant', (1, 5, 1, 14), frozenset({1, 2, 3}), None)])),
+('Expression', ('Constant', (1, 0, 1, 9), (1, 2, 3), None)),
+('Expression', ('Constant', (1, 0, 1, 12), 1, None)),
+]
+eval_results_not_folded =[('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('Add',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('Sub',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('Mult',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('Div',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 6), ('Constant', (1, 0, 1, 1), 1, None), ('FloorDiv',), ('Constant', (1, 5, 1, 6), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('Mod',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 6), ('Constant', (1, 0, 1, 1), 1, None), ('Pow',), ('Constant', (1, 5, 1, 6), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 6), ('Constant', (1, 0, 1, 1), 1, None), ('LShift',), ('Constant', (1, 5, 1, 6), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 6), ('Constant', (1, 0, 1, 1), 1, None), ('RShift',), ('Constant', (1, 5, 1, 6), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('BitOr',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('BitXor',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('BinOp', (1, 0, 1, 5), ('Constant', (1, 0, 1, 1), 1, None), ('BitAnd',), ('Constant', (1, 4, 1, 5), 1, None))),
+('Expression', ('UnaryOp', (1, 0, 1, 2), ('Invert',), ('Constant', (1, 1, 1, 2), 1, None))),
+('Expression', ('UnaryOp', (1, 0, 1, 2), ('USub',), ('Constant', (1, 1, 1, 2), 1, None))),
+('Expression', ('UnaryOp', (1, 0, 1, 2), ('UAdd',), ('Constant', (1, 1, 1, 2), 1, None))),
+('Expression', ('Compare', (1, 0, 1, 14), ('Constant', (1, 0, 1, 1), 1, None), [('In',)], [('List', (1, 5, 1, 14), [('Constant', (1, 6, 1, 7), 1, None), ('Constant', (1, 9, 1, 10), 2, None), ('Constant', (1, 12, 1, 13), 3, None)], ('Load',))])),
+('Expression', ('Compare', (1, 0, 1, 14), ('Constant', (1, 0, 1, 1), 1, None), [('In',)], [('Set', (1, 5, 1, 14), [('Constant', (1, 6, 1, 7), 1, None), ('Constant', (1, 9, 1, 10), 2, None), ('Constant', (1, 12, 1, 13), 3, None)])])),
+('Expression', ('Tuple', (1, 0, 1, 9), [('Constant', (1, 1, 1, 2), 1, None), ('Constant', (1, 4, 1, 5), 2, None), ('Constant', (1, 7, 1, 8), 3, None)], ('Load',))),
+('Expression', ('Subscript', (1, 0, 1, 12), ('Tuple', (1, 0, 1, 9), [('Constant', (1, 1, 1, 2), 1, None), ('Constant', (1, 4, 1, 5), 2, None), ('Constant', (1, 7, 1, 8), 3, None)], ('Load',)), ('Constant', (1, 10, 1, 11), 0, None), ('Load',))),
+]
+exec_results_folded = [('Module', [('For', (1, 0, 2, 5), ('Name', (1, 4, 1, 5), '_', ('Store',)), ('Constant', (1, 9, 1, 18), (1, 2, 3), None), [('Pass', (2, 1, 2, 5))], [], None)], []),
+('Module', [('For', (1, 0, 2, 5), ('Name', (1, 4, 1, 5), '_', ('Store',)), ('Constant', (1, 9, 1, 18), frozenset({1, 2, 3}), None), [('Pass', (2, 1, 2, 5))], [], None)], []),
+]
+exec_results_not_folded =[('Module', [('For', (1, 0, 2, 5), ('Name', (1, 4, 1, 5), '_', ('Store',)), ('List', (1, 9, 1, 18), [('Constant', (1, 10, 1, 11), 1, None), ('Constant', (1, 13, 1, 14), 2, None), ('Constant', (1, 16, 1, 17), 3, None)], ('Load',)), [('Pass', (2, 1, 2, 5))], [], None)], []),
+('Module', [('For', (1, 0, 2, 5), ('Name', (1, 4, 1, 5), '_', ('Store',)), ('Set', (1, 9, 1, 18), [('Constant', (1, 10, 1, 11), 1, None), ('Constant', (1, 13, 1, 14), 2, None), ('Constant', (1, 16, 1, 17), 3, None)]), [('Pass', (2, 1, 2, 5))], [], None)], []),
+]
+
 main()

Lib/test/test_ast/test_ast.py

@@ -21,7 +21,9 @@
 from test.support.ast_helper import ASTTestMixin
 from test.test_ast.utils import to_tuple
 from test.test_ast.snippets import (
-    eval_tests, eval_results, exec_tests, exec_results, single_tests, single_results
+    eval_tests, eval_results, exec_tests, exec_results, single_tests, single_results,
+    eval_opt_tests, eval_results_folded, eval_results_not_folded, exec_opt_tests,
+    exec_results_folded, exec_results_not_folded,
 )
 
 
@@ -135,6 +137,21 @@ def test_optimization_levels_const_folding(self):
                     res = to_tuple(tree.body[0])
                     self.assertEqual(res, expected)
 
+    def test_const_folding(self):
+        for input, folded, not_folded, kind in (
+            (eval_opt_tests, eval_results_folded, eval_results_not_folded, "eval"),
+            (exec_opt_tests, exec_results_folded, exec_results_not_folded, "exec"),
+            ):
+            for i, f, nf in zip(input, folded, not_folded):
+                with self.subTest(action="folding", input=i):
+                    ast_tree = ast.parse(i, "?", kind, optimize=1)
+                    self.assertEqual(to_tuple(ast_tree), f)
+                    self._assertTrueorder(ast_tree, (0, 0))
+                with self.subTest(action="not_folding", input=i):
+                    ast_tree = ast.parse(i, "?", kind)
+                    self.assertEqual(to_tuple(ast_tree), nf)
+                    self._assertTrueorder(ast_tree, (0, 0))
+
     def test_invalid_position_information(self):
         invalid_linenos = [
             (10, 1), (-10, -11), (10, -11), (-5, -2), (-5, 1)

Lib/test/test_ast/utils.py

@@ -3,6 +3,8 @@ def to_tuple(t):
         return t
     elif isinstance(t, list):
         return [to_tuple(e) for e in t]
+    elif isinstance(t, (tuple, frozenset)):
+        return t
     result = [t.__class__.__name__]
     if hasattr(t, 'lineno') and hasattr(t, 'col_offset'):
         result.append((t.lineno, t.col_offset))