Python AST插入print语句未换行问题排查求助

你遇到的问题主要出在两个核心点：AST节点的结构不合法，以及没有将print语句作为独立同级语句插入。咱们一步步拆解解决：

1. 你生成的print节点类型错误 #

在Python AST中，print()是一个表达式，必须包裹在ast.Expr节点里才能成为合法的语句。如果直接返回ast.Call节点，它不属于有效的stmt（语句）类型，astor在生成代码时会把它和for循环挤在同一行——因为它不被视为独立的语句单元。

正确的print语句AST构造应该是这样的：

def get_print_stmt(self):
    # 构造print函数调用的表达式
    print_call = ast.Call(
        func=ast.Name(id='print', ctx=ast.Load()),
        args=[ast.Constant(value="For loop finished!")],  # 替换成你需要打印的内容
        keywords=[]
    )
    # 把表达式包裹成合法的语句节点
    return ast.Expr(value=print_call)

2. 没有将print语句插入到for循环的同级位置 #

你需要把print语句作为独立的同级语句，放在for循环所在的语句块中（比如模块的body、if的body等），而不是修改for循环自身的body（那会把print放到循环体内部，变成每次循环都执行）。

在NodeTransformer中，处理ast.For节点时，你可以返回一个列表，包含原for节点和新的print节点——这样父节点的语句列表会自动把这两个语句按顺序排列：

class CodeInstrumentator(ast.NodeTransformer):
    def get_print_stmt(self):
        print_call = ast.Call(
            func=ast.Name(id='print', ctx=ast.Load()),
            args=[ast.Constant(value="For loop finished!")],
            keywords=[]
        )
        return ast.Expr(value=print_call)

    def visit_For(self, node):
        # 先递归处理for循环内部的子节点
        node = self.generic_visit(node)
        # 返回原for节点 + 新的print语句，替换原来的单个for节点
        return [node, self.get_print_stmt()]

3. 正确处理位置信息 #

现在需要确保AST节点的位置信息正确，让astor能正确排版。处理完整个AST树后，调用ast.fix_missing_locations()补全缺失的位置信息即可——之前单独用increment_lineno()没用，是因为它只对已有行号的节点生效，而你新生成的print节点一开始没有行号。

完整的使用流程示例：

import ast
import astor

# 上面的CodeInstrumentator类代码

# 测试源码
source_code = """
for i in range(3):
    print(i)
for j in range(2):
    print(j*2)
"""

# 解析AST
tree = ast.parse(source_code)
# 转换AST
instrumentator = CodeInstrumentator()
new_tree = instrumentator.visit(tree)
# 补全所有节点的位置信息
ast.fix_missing_locations(new_tree)
# 生成格式化后的代码
print(astor.to_source(new_tree))

运行后生成的代码会是：

for i in range(3):
    print(i)
print('For loop finished!')
for j in range(2):
    print(j * 2)
print('For loop finished!')

为什么之前的方法无效？ #

• fix_missing_locations()是递归补全节点位置，但前提是你的节点本身是合法的语句类型——如果连语句都不是，这个方法也无法帮你排版。
• increment_lineno()是给已有行号的节点增加行号，而你新生成的print节点一开始没有行号，所以调用它完全起不到作用。

内容的提问来源于stack exchange，提问作者Marc Fawver