feat: add ASTTransformer for in-place AST node replacement

luaparser/ast.py

@@ -98,6 +98,82 @@ def to_pretty_json(root: Node) -> str:
     return json.dumps(root, cls=JSONEncoder, indent=4)
 
 
+
+class ASTTransformer:
+    """An iterative depth-first AST visitor that supports in-place node replacement.
+
+    Subclass and override visit_<NodeType>(node) methods. Each method receives
+    the current node and may return:
+
+    - A new :class:`~luaparser.astnodes.Node` instance to replace the current node
+    - ``None`` or the same node to leave it unchanged
+
+    The transformer propagates replacements up to the parent (or returns a new
+    root from :meth:`visit`). After a replacement, children of the *new* node
+    are visited -- not the original.
+
+    Example::
+
+        class NumberDoubler(ASTTransformer):
+            def visit_Number(self, node):
+                return Number(node.value * 2)
+
+        tree = ast.parse("x = 5")
+        new_tree = NumberDoubler().visit(tree)
+    """
+
+    def visit(self, root):
+        """Transform *root* and return the (possibly new) root node.
+
+        Returns None if *root* is None.
+        """
+        if root is None:
+            return None
+
+        # Each stack entry is (node, parent_info) where parent_info is either
+        # None (for root) or (key, container) -- container is a Node or list.
+        node_stack = [(root, None)]
+
+        while node_stack:
+            node, parent_info = node_stack.pop()
+
+            if isinstance(node, Node):
+                # --- call visitor ---
+                name = "visit_" + node.__class__.__name__
+                visitor_method = getattr(self, name, None)
+                if visitor_method is not None:
+                    replacement = visitor_method(node)
+                    if replacement is not None and replacement is not node:
+                        # Replace node in parent (or update root)
+                        if parent_info is not None:
+                            parent_key, parent_container = parent_info
+                            if isinstance(parent_container, list):
+                                parent_container[parent_key] = replacement
+                            else:
+                                setattr(parent_container, parent_key, replacement)
+                        else:
+                            root = replacement
+                        node = replacement  # visit replacement's children
+
+                # --- push children (reverse order for correct DFS) ---
+                children = [
+                    attr for attr in node.__dict__.keys()
+                    if not attr.startswith("_")
+                ]
+                for child_key in reversed(children):
+                    child = node.__dict__[child_key]
+                    if isinstance(child, list):
+                        for i in reversed(range(len(child))):
+                            node_stack.append((child[i], (i, child)))
+                    elif isinstance(child, Node):
+                        node_stack.append((child, (child_key, node)))
+
+            elif isinstance(node, list):
+                for n in reversed(node):
+                    node_stack.append((n, parent_info))
+
+        return root
+
 class ASTVisitor:
     def visit(self, root):
         # base case:

luaparser/tests/test_ast.py

@@ -158,3 +158,251 @@ def test_to_pretty_json(self):
             }"""
         )
         self.assertEqual(ast.to_pretty_json(ast.parse(src)), exp)
+
+
+class ASTTransformerTestCase(tests.TestCase):
+    """Tests for ASTTransformer."""
+
+    def test_noop_returns_same_tree(self):
+        """Transformer with no overrides returns identical structure."""
+        src = "local x = 1"
+        tree = ast.parse(src)
+
+        class NoopTransformer(ast.ASTTransformer):
+            pass
+
+        result = NoopTransformer().visit(tree)
+        self.assertEqual(ast.to_pretty_str(result), ast.to_pretty_str(tree))
+        # Root should be the same object when no transforms applied
+        self.assertIs(result, tree)
+
+    def test_replace_number(self):
+        """Replace a Number node with a different value."""
+        src = "local x = 42"
+        tree = ast.parse(src)
+
+        class Doubler(ast.ASTTransformer):
+            def visit_Number(self, node):
+                return Number(node.n * 2)
+
+        result = Doubler().visit(tree)
+        self.assertIn("84", ast.to_pretty_str(result))
+        self.assertNotIn("42", ast.to_pretty_str(result))
+
+    def test_replace_string(self):
+        """Replace a String literal."""
+        src = 'local msg = "hello"'
+        tree = ast.parse(src)
+
+        class StringUpper(ast.ASTTransformer):
+            def visit_String(self, node):
+                upper_raw = node.raw.upper()
+                return String(upper_raw.encode(), upper_raw, node.delimiter)
+
+        result = StringUpper().visit(tree)
+        output = ast.to_pretty_str(result)
+        self.assertIn("HELLO", output)
+        self.assertNotIn("hello", output)
+
+    def test_keep_node_by_returning_none(self):
+        """Returning None keeps the original node."""
+        src = "local x = 42"
+        tree = ast.parse(src)
+
+        class KeepNumbers(ast.ASTTransformer):
+            def visit_Number(self, node):
+                if node.n == 42:
+                    return None  # keep
+                return Number(0)
+
+        result = KeepNumbers().visit(tree)
+        self.assertEqual(ast.to_pretty_str(result), ast.to_pretty_str(tree))
+
+    def test_keep_node_by_returning_same(self):
+        """Returning the same node object keeps it."""
+        src = "local x = 42"
+        tree = ast.parse(src)
+
+        class IdentityTransformer(ast.ASTTransformer):
+            def visit_Number(self, node):
+                return node  # same object
+
+        result = IdentityTransformer().visit(tree)
+        self.assertEqual(ast.to_pretty_str(result), ast.to_pretty_str(tree))
+        self.assertIs(result, tree)
+
+    def test_replace_statement(self):
+        """Replace an entire statement node."""
+        src = "local x = 1"
+        tree = ast.parse(src)
+
+        class AssignToNil(ast.ASTTransformer):
+            def visit_LocalAssign(self, node):
+                # Replace local assignment with a nil assignment
+                return Assign(node.targets, [Nil()])
+
+        result = AssignToNil().visit(tree)
+        output = ast.to_lua_source(result)
+        self.assertIn("nil", output)
+        self.assertNotIn("local", output)
+
+    def test_replace_multiple_in_list(self):
+        """Replace multiple nodes in a list (block body)."""
+        src = textwrap.dedent("""\
+            local a = 1
+            local b = 2
+            local c = 3
+        """)
+        tree = ast.parse(src)
+
+        class DropMiddleStatement(ast.ASTTransformer):
+            def visit_LocalAssign(self, node):
+                target_name = node.targets[0].id
+                if target_name == "b":
+                    return None  # remove it -- handled via parent removal
+                return node
+
+        result = DropMiddleStatement().visit(tree)
+        output = ast.to_lua_source(result)
+        self.assertIn("local a", output)
+        self.assertIn("local c", output)
+        # 'local b' should still be there since returning None keeps the node
+        self.assertIn("local b", output)
+
+    def test_replace_in_call_args(self):
+        """Replace an argument inside a function call."""
+        src = "print(42)"
+        tree = ast.parse(src)
+
+        class ArgReplacer(ast.ASTTransformer):
+            def visit_Number(self, node):
+                s = str(node.n).encode()
+                return String(s, str(node.n), StringDelimiter.DOUBLE_QUOTE)
+
+        result = ArgReplacer().visit(tree)
+        output = ast.to_pretty_str(result)
+        self.assertIn("String", output)
+
+    def test_recursive_replacement(self):
+        """Perform recursive transformations (expression folding)."""
+        src = "local x = 2 + 3"
+        tree = ast.parse(src)
+
+        class ConstantFolder(ast.ASTTransformer):
+            def visit_AddOp(self, node):
+                if isinstance(node.left, Number) and isinstance(node.right, Number):
+                    return Number(node.left.n + node.right.n)
+                return node
+
+        result = ConstantFolder().visit(tree)
+        output = ast.to_lua_source(result)
+        self.assertIn("= 5", output)
+
+    def test_chained_replacement(self):
+        """Replace a node, then replace a child of the replacement."""
+        src = textwrap.dedent("""\
+            if true then
+                x = 1
+            end
+        """)
+        tree = ast.parse(src)
+
+        class IfRewriter(ast.ASTTransformer):
+            def visit_If(self, node):
+                # Replace all if-statements with a do-block (for testing)
+                return Do(node.body)
+            def visit_TrueExpr(self, node):
+                # Replace true with false (tests ordering: If replaced first,
+                # so TrueExpr inside old If condition is never visited)
+                return FalseExpr()
+
+        result = IfRewriter().visit(tree)
+        output = ast.to_lua_source(result)
+        self.assertIn("do", output.lower())
+
+    def test_none_root_returns_none(self):
+        """Passing None returns None."""
+        class NullTransformer(ast.ASTTransformer):
+            pass
+        self.assertIsNone(NullTransformer().visit(None))
+
+    def test_root_replacement(self):
+        """Replace the root Chunk node."""
+        src = "local x = 1"
+        tree = ast.parse(src)
+
+        class RootReplacer(ast.ASTTransformer):
+            def visit_Chunk(self, node):
+                new_body = Block([LocalAssign(
+                    [Name("y")],
+                    [Number(99)]
+                )])
+                return Chunk(new_body)
+
+        result = RootReplacer().visit(tree)
+        output = ast.to_lua_source(result)
+        self.assertIn("y = 99", output)
+        self.assertNotIn("x = 1", output)
+
+    def test_identity_transformer_roundtrip(self):
+        """Full round-trip: parse -> transform (identity) -> generate Lua."""
+        src = textwrap.dedent("""\
+            function foo(a, b)
+                if a > b then
+                    return a + 1
+                else
+                    return b * 2
+                end
+            end
+            foo(10, 20)
+        """)
+        tree = ast.parse(src)
+        original_output = ast.to_lua_source(tree)
+
+        class Identity(ast.ASTTransformer):
+            pass
+
+        result = Identity().visit(tree)
+        result_output = ast.to_lua_source(result)
+
+        self.assertEqual(original_output, result_output)
+        self.assertIs(result, tree)
+
+    def test_visit_method_for_parent_class(self):
+        """visit_Name is called for Name nodes; no crash on unrecognized types."""
+        src = "local x = 1"
+        tree = ast.parse(src)
+
+        visited_names = []
+
+        class NameCollector(ast.ASTTransformer):
+            def visit_Name(self, node):
+                visited_names.append(node.id)
+                return node
+
+        NameCollector().visit(tree)
+        self.assertIn("x", visited_names)
+
+    def test_replace_deeply_nested(self):
+        """Replace a node deep in the tree."""
+        src = textwrap.dedent("""\
+            local t = {
+                a = {
+                    b = 42
+                }
+            }
+        """)
+        tree = ast.parse(src)
+
+        class DeepReplacer(ast.ASTTransformer):
+            def visit_Number(self, node):
+                if node.n == 42:
+                    return Number(999)
+                return node
+
+        result = DeepReplacer().visit(tree)
+        output = ast.to_lua_source(result)
+        self.assertIn("999", output)
+        self.assertNotIn("42", output)
+
+