Python标准库中的ast模块一直没怎么使用,这个过于底层的模块在一般逻辑中很少涉及,刚好最近有个改动,清理代码中残存的print语句,需要通过它来完成。

针对这个问题最直接的想法是直接对代码文件通过正则进行匹配,查找print开头的代码行进行替换。不过粗暴的正则匹配只能解决单行问题,对于跨行的语句是无能为力的。为了保证正确性,需要另外想办法。在一番搜索之后找到的方法是通过标准库中的ast模块将代码解析成AST之后,遍历AST,替换掉print语句节点,再将修改后的语法树导出成代码文件。步骤非常简单,实现也并不复杂。

AST的解析与导出

StackOverflow的这个问题的答案中推荐了astor用于将AST导出成代码文件。代码解析可以使用此模块也可以直接使用ast。

def parse_and_export(input_path, output_path):
    with open(output_path, 'w') as outputs:
        outputs.write(astor.to_source(astor.parsefile(input_path)))
    outputs.flush()

替换AST节点

上面parse_and_export只是进行了简单的解析与导出并未对AST进行修改,那么如何进行修改替换?

ast文档中对此进行了介绍,

class ast.NodeTransformer
    A NodeVisitor subclass that walks the abstract syntax tree and allows modification of nodes.
    The NodeTransformer will walk the AST and use the return value of the visitor methods to replace or remove the old node. If the return value of the visitor method is None, the node will be removed from its location, otherwise it is replaced with the return value. The return value may be the original node in which case no replacement takes place.

由此可知,继承NodeTransformer,并提供相应的visit_xxx方法即能够完成需求。

class RemovePrint(NodeTransformer):

    def visit_Print(self, node):
        return ast.Pass()

将上述代码整合,

def remove_print(input_path, output_path):
    with open(output_path, 'w') as outputs:
        outputs.write(astor.to_source(RemovePrint().visit(astor.parsefile(input_path))))
    outputs.flush()

参考