Adds ast rewriter

DinoV · DinoV · commit 90e344988361 · 2021-01-28T11:14:38.000-08:00
Summary:
Python 3.7 starts doing various optimizations at the AST level, where it actually makes sense.  To do this we need the ability to re-write the ASTs.  This adds an ASTRewriter class which supports having the visit() methods return the new nodes to be re-written.  Child nodes which aren't modified remain in the tree, and the spine of the AST is rewritten as we walk back up the AST.

Even though the Python ASTs are mutable this supports performing the re-writes w/o modifying the AST in place.

Test Plan: ./python -m test.test_compiler
diff --git a/compiler/visitor.py b/compiler/visitor.py
@@ -1,6 +1,7 @@
 from __future__ import print_function
 
 import ast
+from ast import AST, copy_location
 from typing import Any, List
 
 # XXX should probably rename ASTVisitor to ASTWalker
@@ -63,6 +64,75 @@ def visit(self, node, *args):
         return meth(node, *args)
 
 
+class ASTRewriter(ASTVisitor):
+    """performs rewrites on the AST, rewriting parent nodes when child nodes
+    are replaced."""
+
+    @staticmethod
+    def update_node(node: AST, **replacement: Any) -> AST:
+        res = node
+        for name, val in replacement.items():
+            existing = getattr(res, name)
+            if existing is val:
+                continue
+
+            if node is res:
+                res = ASTRewriter.clone_node(node)
+
+            setattr(res, name, val)
+        return res
+
+    @staticmethod
+    def clone_node(node: AST) -> AST:
+        attrs = []
+        for name in node._fields:
+            attr = getattr(node, name, None)
+            if isinstance(attr, list):
+                attr = list(attr)
+            attrs.append(attr)
+
+        new = type(node)(*attrs)
+        return copy_location(new, node)
+
+    def walk_list(self, old_values: List[AST]) -> List[AST]:
+        new_values = []
+        changed = False
+        for value in old_values:
+            if isinstance(value, AST):
+                new_value = self.visit(value)
+                changed |= new_value is not value
+                if new_value is None:
+                    continue
+                elif not isinstance(new_value, AST):
+                    new_values.extend(new_value)
+                    continue
+                value = new_value
+
+            new_values.append(value)
+        return new_values if changed else old_values
+
+    def generic_visit(self, node: AST, *args) -> AST:
+        if isinstance(node, list):
+            return self.walk_list(node)
+
+        ret_node = node
+        for field, old_value in ast.iter_fields(node):
+            if not isinstance(old_value, (AST, list)):
+                continue
+
+            new_node = self.visit(old_value)
+            assert (  # noqa: IG01
+                new_node is not None
+            ), "can't remove AST nodes that aren't part of a list"
+            if new_node is not old_value:
+                if ret_node is node:
+                    ret_node = self.clone_node(node)
+
+                setattr(ret_node, field, new_node)
+
+        return ret_node
+
+
 class ExampleASTVisitor(ASTVisitor):
     """Prints examples of the nodes that aren't visited
 
diff --git a/test_compiler/__init__.py b/test_compiler/__init__.py
@@ -8,3 +8,4 @@
 from .test_sbs_stdlib import SbsCompileTests
 from .test_symbols import SymbolVisitorTests
 from .test_unparse import UnparseTests
+from .test_visitor import VisitorTests
diff --git a/test_compiler/__main__.py b/test_compiler/__main__.py
@@ -9,6 +9,7 @@
 from .test_errors import ErrorTests, ErrorTestsBuiltin
 from .test_symbols import SymbolVisitorTests
 from .test_unparse import UnparseTests
+from .test_visitor import VisitorTests
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test_compiler/test_visitor.py b/test_compiler/test_visitor.py
@@ -0,0 +1,88 @@
+import ast
+from compiler.visitor import ASTRewriter
+from compiler.unparse import to_expr
+from unittest import TestCase
+
+
+class VisitorTests(TestCase):
+    def test_rewrite_node(self):
+        class TestRewriter(ASTRewriter):
+            def visitName(self, node):
+                if isinstance(node, ast.Name) and node.id == "z":
+                    return ast.Name("foo", ast.Load())
+                return node
+
+        tree = ast.parse("x + y + z").body[0].value
+
+        rewriter = TestRewriter()
+        new_tree = rewriter.visit(tree)
+
+        self.assertIsNotNone(new_tree)
+        self.assertNotEqual(new_tree, tree, "tree should be rewritten")
+        self.assertEqual(to_expr(new_tree), "x + y + foo")
+        self.assertEqual(tree.left, new_tree.left, "Unchanged nodes should be the same")
+
+    def test_rewrite_stmt(self):
+        class TestRewriter(ASTRewriter):
+            def visitAssign(self, node):
+                return ast.AnnAssign(
+                    ast.Name("foo", ast.Store), ast.Str("foo"), ast.Num(1), True
+                )
+
+        tree = ast.parse("x = 1\nf()\n")
+
+        rewriter = TestRewriter()
+        new_tree = rewriter.visit(tree)
+
+        self.assertIsNotNone(new_tree)
+        self.assertNotEqual(new_tree, tree, "tree should be rewritten")
+        self.assertNotEqual(tree.body[0], new_tree.body[0])
+        self.assertEqual(tree.body[1], new_tree.body[1])
+
+    def test_remove_node(self):
+        class TestRewriter(ASTRewriter):
+            def visitAssign(self, node):
+                return None
+
+        tree = ast.parse("x = 1\nf()\n")
+
+        rewriter = TestRewriter()
+        new_tree = rewriter.visit(tree)
+
+        self.assertIsNotNone(new_tree)
+        self.assertEqual(tree.body[1], new_tree.body[0])
+        self.assertEqual(len(new_tree.body), 1)
+
+    def test_change_child_and_list(self):
+        class TestRewriter(ASTRewriter):
+            def visitarguments(self, node: ast.arguments):
+                node = self.clone_node(node)
+                node.vararg = ast.arg("args", None)
+                return node
+
+            def visitAssign(self, node):
+                return ast.Pass()
+
+        tree = ast.parse("def f():\n    x = 1")
+
+        rewriter = TestRewriter()
+        new_tree = rewriter.visit(tree)
+        func = new_tree.body[0]
+        self.assertEqual(type(func.body[0]), ast.Pass)
+        self.assertIsNotNone(func.args.vararg)
+
+    def test_change_list_and_child(self):
+        class TestRewriter(ASTRewriter):
+            def visitStr(self, node: ast.Str):
+                return ast.Str("bar")
+
+            def visitAssign(self, node):
+                return ast.Pass()
+
+        tree = ast.parse("def f() -> 'foo':\n    x = 1")
+
+        rewriter = TestRewriter()
+        new_tree = rewriter.visit(tree)
+        func = new_tree.body[0]
+        self.assertIsNotNone(func.returns)
+        self.assertEqual(type(func.body[0]), ast.Pass)
