我想以编程方式编辑python源代码。基本上我想读取一个.py文件,生成AST,然后写回修改后的python源代码(即另一个.py文件)。

有一些方法可以使用标准的python模块(如ast或compiler)来解析/编译python源代码。但是,我认为它们都不支持修改源代码的方法(例如删除这个函数声明),然后再写回修改的python源代码。

更新:我想这样做的原因是我想为python写一个突变测试库,主要是通过删除语句/表达式,重新运行测试,看看有什么破坏。


当前回答

我最近创建了相当稳定的(核心是经过良好测试的)和可扩展的代码段,它从ast树生成代码:https://github.com/paluh/code-formatter。

我正在使用我的项目作为一个小vim插件的基础(我每天都在使用),所以我的目标是生成非常漂亮和可读的python代码。

P.S. I've tried to extend codegen but it's architecture is based on ast.NodeVisitor interface, so formatters (visitor_ methods) are just functions. I've found this structure quite limiting and hard to optimize (in case of long and nested expressions it's easier to keep objects tree and cache some partial results - in other way you can hit exponential complexity if you want to search for best layout). BUT codegen as every piece of mitsuhiko's work (which I've read) is very well written and concise.

其他回答

我们也有类似的需求,这里的其他答案并没有解决这个问题。因此,我们为此创建了一个库ASTTokens,它采用AST或astroid模块生成的AST树,并用原始源代码中的文本范围标记它。

它不直接修改代码,但在上面添加代码并不难,因为它会告诉您需要修改的文本范围。

例如,这将在WRAP(…)中包装一个函数调用,保留注释和其他内容:

example = """
def foo(): # Test
  '''My func'''
  log("hello world")  # Print
"""

import ast, asttokens
atok = asttokens.ASTTokens(example, parse=True)

call = next(n for n in ast.walk(atok.tree) if isinstance(n, ast.Call))
start, end = atok.get_text_range(call)
print(atok.text[:start] + ('WRAP(%s)' % atok.text[start:end])  + atok.text[end:])

生产:

def foo(): # Test
  '''My func'''
  WRAP(log("hello world"))  # Print

希望这能有所帮助!

不幸的是,上面的答案实际上没有一个同时满足这两个条件

保持周围源代码的语法完整性(例如保留注释,其他类型的代码格式) 实际上使用AST(而不是CST)。

我最近写了一个小工具包来进行纯基于AST的重构,称为重构。例如,如果你想用42替换所有占位符,你可以简单地像这样写一个规则;

class Replace(Rule):
    
    def match(self, node):
        assert isinstance(node, ast.Name)
        assert node.id == 'placeholder'
        
        replacement = ast.Constant(42)
        return ReplacementAction(node, replacement)

它会找到所有可接受的节点,用新节点替换它们并生成最终的表单;

--- test_file.py
+++ test_file.py

@@ -1,11 +1,11 @@

 def main():
-    print(placeholder * 3 + 2)
-    print(2 +               placeholder      + 3)
+    print(42 * 3 + 2)
+    print(2 +               42      + 3)
     # some commments
-    placeholder # maybe other comments
+    42 # maybe other comments
     if something:
         other_thing
-    print(placeholder)
+    print(42)
 
 if __name__ == "__main__":
     main()

Pythoscope对它自动生成的测试用例执行此操作,就像python 2.6的2to3工具一样(它转换python 2。X源代码转换为python 3。x源)。

这两个工具都使用lib2to3库,它是python解析器/编译器机制的实现,可以在从source -> AST -> source循环绊倒源代码时保留源代码中的注释。

如果您想进行更多的重构(如转换),rope项目可以满足您的需求。

ast模块是另一个选择,还有一个关于如何将语法树“解解析”回代码的旧示例(使用解析器模块)。但是ast模块在对随后转换为代码对象的代码进行ast转换时更有用。

红男爵计划也可能是个不错的选择(泽维尔·康贝尔)

在ast模块的帮助下,解析和修改代码结构当然是可能的,我将在稍后的示例中展示它。然而,仅使用ast模块是不可能写回修改后的源代码的。还有其他模块可用于这项工作,例如这里的一个。

注意:下面的例子可以作为ast模块使用的入门教程,但是更全面的ast模块使用指南可以在绿树蛇教程和ast模块的官方文档中找到。

ast简介:

>>> import ast
>>> tree = ast.parse("print 'Hello Python!!'")
>>> exec(compile(tree, filename="<ast>", mode="exec"))
Hello Python!!

你可以通过调用API ast.parse()来解析python代码(以字符串表示)。它返回抽象语法树(AST)结构的句柄。有趣的是,您可以编译回这个结构并执行它,如上面所示。

另一个非常有用的API是AST .dump(),它将整个AST以字符串形式转储。它可以用来检查树形结构,在调试中有很大的帮助。例如,

在Python 2.7中:

>>> import ast
>>> tree = ast.parse("print 'Hello Python!!'")
>>> ast.dump(tree)
"Module(body=[Print(dest=None, values=[Str(s='Hello Python!!')], nl=True)])"

在Python 3.5上:

>>> import ast
>>> tree = ast.parse("print ('Hello Python!!')")
>>> ast.dump(tree)
"Module(body=[Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Str(s='Hello Python!!')], keywords=[]))])"

请注意Python 2.7和Python 3.5中print语句的语法差异,以及各自树中AST节点类型的差异。


如何使用ast修改代码:

现在,让我们看一个用ast模块修改python代码的例子。修改AST结构的主要工具是AST . nodetransformer类。每当一个人需要修改AST时,他/她需要从它继承子类并相应地编写节点转换。

对于我们的例子,让我们试着编写一个简单的实用程序,它可以将python2, print语句转换为python3函数调用。

打印语句到Fun调用转换工具:print2to3.py

#!/usr/bin/env python
'''
This utility converts the python (2.7) statements to Python 3 alike function calls before running the code.

USAGE:
     python print2to3.py <filename>
'''
import ast
import sys

class P2to3(ast.NodeTransformer):
    def visit_Print(self, node):
        new_node = ast.Expr(value=ast.Call(func=ast.Name(id='print', ctx=ast.Load()),
            args=node.values,
            keywords=[], starargs=None, kwargs=None))
        ast.copy_location(new_node, node)
        return new_node

def main(filename=None):
    if not filename:
        return

    with open(filename, 'r') as fp:
        data = fp.readlines()
    data = ''.join(data)
    tree = ast.parse(data)

    print "Converting python 2 print statements to Python 3 function calls"
    print "-" * 35
    P2to3().visit(tree)
    ast.fix_missing_locations(tree)
    # print ast.dump(tree)

    exec(compile(tree, filename="p23", mode="exec"))

if __name__ == '__main__':
    if len(sys.argv) <=1:
        print ("\nUSAGE:\n\t print2to3.py <filename>")
        sys.exit(1)
    else:
        main(sys.argv[1])

这个实用程序可以在一个小的示例文件上尝试,比如下面的一个,它应该可以正常工作。

输入文件:py2.py

class A(object):
    def __init__(self):
        pass

def good():
    print "I am good"

main = good

if __name__ == '__main__':
    print "I am in main"
    main()

请注意,上面的转换仅供上一个教程使用,在实际情况下,必须查看所有不同的场景,例如打印“x是%s”%(“Hello Python”)。

内置ast模块似乎没有转换回源代码的方法。但是,这里的codegen模块为ast提供了一个漂亮的打印机,使您能够这样做。 如。

import ast
import codegen

expr="""
def foo():
   print("hello world")
"""
p=ast.parse(expr)

p.body[0].body = [ ast.parse("return 42").body[0] ] # Replace function body with "return 42"

print(codegen.to_source(p))

这将打印:

def foo():
    return 42

请注意,您可能会丢失确切的格式和注释,因为这些没有保留。

但是,您可能不需要这样做。如果您所需要的只是执行替换的AST,那么只需在AST上调用compile()并执行结果代码对象即可。