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)
+
+

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)
+
+

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)
+
+

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)
+
+

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)
+
+

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)
+
+

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)
+
+

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)
+
+