libcst_function_rename_transformer_preserving_formatting.py

1import libcst as cst
2
3class TypingCollector(cst.CSTTransformer):
4    def leave_FunctionDef(
5        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
6    ) -> cst.FunctionDef:
7        if original_node.name.value == "bad_name":
8            return updated_node.with_changes(name=cst.Name("good_name"))
9        return updated_node
10
11# Input source code
12source_code = """
13def bad_name():
14    pass
15"""
16
17# Parse the code into a CST
18source_tree = cst.parse_module(source_code)
19
20# Transform the tree
21transformer = TypingCollector()
22modified_tree = source_tree.visit(transformer)
23
24# Print the modified code
25print(modified_tree.code)