google_pasta_format_preserving_ast_parse_modify_and_dump.py

1import ast
2import pasta
3
4# 1. Provide the source code to be refactored
5source_code = """
6def hello_world():
7  # This comment will be preserved
8  print("Hello, world!")
9"""
10
11# 2. Parse the source code using pasta
12# This creates an AST that tracks formatting information
13tree = pasta.parse(source_code)
14
15# 3. Perform a modification using standard 'ast' nodes
16# For example, let's change the function name
17for node in ast.walk(tree):
18    if isinstance(node, ast.FunctionDef) and node.name == 'hello_world':
19        node.name = 'greet'
20
21# 4. Convert the AST back to code
22# The output will preserve the original comments and indentation
23modified_code = pasta.dump(tree)
24
25print(modified_code)