diff --git a/luaparser/ast.py b/luaparser/ast.py index 647a89b..6d62623 100644 --- a/luaparser/ast.py +++ b/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_(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: diff --git a/luaparser/tests/test_ast.py b/luaparser/tests/test_ast.py index 7c9a3d3..5746552 100644 --- a/luaparser/tests/test_ast.py +++ b/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) + +