解析 .py 文件，读取 AST，修改它，然后写回修改后的源代码

Posted 2023-02-19

【中文标题】解析 .py 文件，读取 AST，修改它，然后写回修改后的源代码

【英文标题】：Parse a .py file, read the AST, modify it, then write back the modified source code

【发布时间】：2010-10-20 14:20:41

【问题描述】：

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

有一些方法可以使用标准的 Python 模块解析/编译 Python 源代码，例如 ast 或 compiler。但是，我认为它们中的任何一个都不支持修改源代码的方式（例如删除此函数声明），然后将修改后的 python 源代码写回。

更新：我想这样做的原因是我想为 python 写一个Mutation testing library，主要是通过删除语句/表达式，重新运行测试并查看什么中断。

【问题讨论】：

自 2.6 版起已弃用：编译器包已在 Python 3.0 中删除。

什么不能编辑源？为什么不能写装饰器？

天哪！我想用相同的技术（特别是创建一个鼻子插件）为 python 做一个突变测试器，你打算开源它吗？

@Ryan 是的，我将开源我创建的任何东西。我们应该就此保持联系

当然，我通过 Launchpad 向您发送了一封电子邮件。

【参考方案1】：

Pythoscope 对它自动生成的测试用例执行此操作，就像用于 python 2.6 的 2to3 工具一样（它将 python 2.x 源代码转换为 python 3.x 源代码）。

这两个工具都使用 lib2to3 库，它是 python 解析器/编译器机制的实现，当它从源 -> AST -> 源往返时，可以在源中保留 cmets。

rope project 可以满足您的需求，如果您想进行更多重构，例如转换。

ast 模块是您的另一个选择，there's an older example of how to "unparse" syntax trees back into code（使用解析器模块）。但是ast 模块在对代码进行 AST 转换然后转换为代码对象时更有用。

redbaron 项目也可能是一个不错的选择 (ht Xavier Combelle)

【讨论】：

unparse 的例子还在维护，这里是更新的 py3k 版本：hg.python.org/cpython/log/tip/Tools/parser/unparse.py

关于unparse.py 脚本 - 从另一个脚本中使用它可能真的很麻烦。但是，有一个名为 astunparse (on github, on pypi) 的包，它基本上是 unparse.py 的正确打包版本。

您能否通过添加 parso 作为首选选项来更新您的答案？非常好，而且更新了。

@Ryan。你能给我一些工具来获取 python 源代码的 AST 和 CFG 吗？

【参考方案2】：

内置的 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

请注意，您可能会丢失确切的格式和 cmets，因为它们不会被保留。

但是，您可能不需要。如果您只需要执行替换的 AST，您只需在 ast 上调用 compile() 并执行生成的代码对象即可。

【讨论】：

对于任何将来使用它的人来说，codegen 基本上已经过时并且有一些错误。我已经修复了其中的几个；我在 github 上有一个要点：gist.github.com/791312

请注意，最新的 codegen 更新于 2012 年，在上述评论之后，所以我猜 codegen 已更新。 @mattbasta

astor 似乎是 codegen 的维护继任者

【参考方案3】：

在另一个答案中，我建议使用 astor 包，但后来我发现了一个更新的 AST 反解析包，名为 astunparse：

>>> import ast
>>> import astunparse
>>> print(astunparse.unparse(ast.parse('def foo(x): return 2 * x')))


def foo(x):
    return (2 * x)

我已经在 Python 3.5 上对此进行了测试。

【参考方案4】：

您可能不需要重新生成源代码。当然，这对我来说有点危险，因为您实际上并没有解释为什么您认为需要生成一个充满代码的 .py 文件；但是：

如果您想生成一个人们将实际使用的 .py 文件，也许这样他们就可以填写表格并获得一个有用的 .py 文件插入到他们的项目中，那么您不想将其更改为 AST 并返回，因为您将丢失 所有格式（想想那些通过将相关的行组合在一起使 Python 如此可读的空行） (ast nodes have lineno and col_offset attributes) cmets。相反，您可能希望使用模板引擎（例如，Django template language 旨在使模板甚至文本文件变得容易）来自定义 .py 文件，或者使用 Rick Copeland 的 MetaPython 扩展名。

如果您在编译模块期间尝试进行更改，请注意您不必一直回到文本；您可以直接编译 AST 而不是将其转回 .py 文件。

1234563如果您扩展您的问题以让我们知道您真正想要完成什么，那么新的 .py 文件可能根本不会参与答案；我已经看到数百个 Python 项目在做数百个现实世界的事情，而且没有一个项目需要编写 .py 文件。所以，我必须承认，我有点怀疑你是否找到了第一个好的用例。 :-)

更新：既然你已经解释了你想要做什么，我还是很想只对 AST 进行操作。您将希望通过删除而不是文件的行来进行变异（这可能导致半语句简单地因 SyntaxError 而死），而是整个语句 - 还有什么比在 AST 中更好的地方来做到这一点？

【讨论】：

很好地概述了可能的解决方案和可能的替代方案。

代码生成的真实世界用例：Kid 和 Genshi（我相信）从 XML 模板生成 Python 以快速呈现动态页面。

【参考方案5】：

花了一些时间，但 Python 3.9 有这个： https://docs.python.org/3.9/whatsnew/3.9.html#ast https://docs.python.org/3.9/library/ast.html#ast.unparse

ast.unparse(ast_obj)

解析一个 ast.AST 对象并生成一个带有代码的字符串，如果用 ast.parse() 解析回来，该代码将产生一个等效的 ast.AST 对象。

【参考方案6】：

在ast 模块的帮助下，解析和修改代码结构当然是可能的，稍后我将在一个示例中展示它。但是，仅使用 ast 模块无法写回修改后的源代码。此作业还有其他可用模块，例如一个here。

注意：下面的示例可以被视为有关使用 ast 模块的介绍性教程，但在 Green Tree snakes tutorial 和 official documentation on ast module 上提供了有关使用 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 时，他/她需要从中继承子类并相应地编写节点转换。

对于我们的示例，让我们尝试编写一个简单的实用程序，将 Python 2 的 print 语句转换为 Python 3 的函数调用。

将语句打印到 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()

请注意，上述转换仅用于 ast 教程目的，在实际情况下，您必须查看所有不同的场景，例如 print " x is %s" % ("Hello Python")。

【讨论】：

这不显示如何打印，它执行？

【参考方案7】：

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

我将我的项目用作小型 vim 插件（我每天都在使用）的基础，所以我的目标是生成非常漂亮且可读的 Python 代码。

附： 我尝试扩展codegen，但它的架构基于ast.NodeVisitor 接口，所以格式化程序（visitor_ 方法）只是函数。我发现这种结构非常有限且难以优化（在长且嵌套的表达式的情况下，保留对象树并缓存一些部分结果更容易 - 以其他方式，如果您想搜索最佳布局，您可以达到指数复杂性）。 但是 codegen mitsuhiko 的每一部作品（我读过）都写得非常好，简洁。

【参考方案8】：

如果你在 2019 年看到这个，那么你可以使用这个libcst 包裹。它的语法类似于 ast。这就像一个魅力，并保留了代码结构。对于必须保留 cmets、空格、换行符等的项目，它基本上很有帮助。

如果您不需要关心保留 cmets、空白和其他内容，那么 ast 和 astor 的组合效果很好。

【参考方案9】：

One of the other answers 推荐codegen，它似乎已被astor 取代。 astor on PyPI 的版本（撰写本文时为 0.5 版）似乎也有点过时，因此您可以安装 astor 的开发版本，如下所示。

pip install git+https://github.com/berkerpeksag/astor.git#egg=astor

然后您可以使用astor.to_source 将 Python AST 转换为人类可读的 Python 源代码：

>>> import ast
>>> import astor
>>> print(astor.to_source(ast.parse('def foo(x): return 2 * x')))
def foo(x):
    return 2 * x

我已经在 Python 3.5 上对此进行了测试。

【参考方案10】：

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

它不会直接对代码进行修改，但这并不难在上面添加，因为它确实告诉您需要修改的文本范围。

例如，这在WRAP(...) 中包装了一个函数调用，保留了 cmets 和其他所有内容：

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

希望这会有所帮助！

【参考方案11】：

很遗憾，上面的答案都没有真正满足这两个条件

保持周围源代码的句法完整性（例如，保留 cmets，其余代码的其他类型的格式） 实际使用 AST（不是 CST）。

我最近编写了一个小工具包来进行纯 AST 的重构，称为 refactor。比如你想把所有的placeholders替换成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()

【参考方案12】：

Program Transformation System 是一种工具，它可以解析源文本、构建 AST，允许您使用源到源的转换来修改它们（“如果你看到这种模式，用那个模式替换它”）。此类工具非常适合对现有源代码进行突变，即“如果您看到这种模式，请用模式变体替换”。

当然，您需要一个程序转换引擎，它可以解析您感兴趣的语言，并且仍然进行模式导向的转换。我们的DMS Software Reengineering Toolkit 是一个可以做到这一点的系统，它可以处理 Python 和各种其他语言。

准确地看到这个SO answer for an example of a DMS-parsed AST for Python capturing comments。 DMS 可以更改 AST，并重新生成有效文本，包括 cmets。您可以要求它使用自己的格式约定（您可以更改这些约定）来漂亮地打印 AST，或者执行“保真打印”，它使用原始行和列信息来最大程度地保留原始布局（新代码的布局中的一些更改插入是不可避免的）。

要使用 DMS 为 Python 实现“变异”规则，您可以编写以下代码：

rule mutate_addition(s:sum, p:product):sum->sum =
  " \s + \p " -> " \s - \p"
 if mutate_this_place(s);

此规则以语法正确的方式将“+”替换为“-”；它在 AST 上运行，因此不会触及碰巧看起来正确的字符串或 cmets。 “mutate_this_place”的额外条件是让您控制这种情况发生的频率；你不想改变程序中的每个地方。

您显然需要更多这样的规则来检测各种代码结构，并用变异版本替换它们。 DMS 很乐意应用一组规则。然后将变异的 AST 打印出来。

【讨论】：

我已经 4 年没看过这个答案了。哇，它已被多次否决。这真是太棒了，因为它直接回答了 OP 的问题，甚至展示了如何做他想做的突变。我不认为任何投反对票的人会愿意解释为什么他们投反对票。

因为它推广了一种非常昂贵的闭源工具。

@ZoranPavlovic：所以你不反对它的任何技术准确性或实用性？

@Zoran：他没有说他有一个开源库。他说他想修改 Python 源代码（使用 AST），但他能找到的解决方案并没有做到这一点。这是一个这样的解决方案。您认为人们不会在使用 Python on Java 等语言编写的程序上使用商业工具吗？

我不是一个反对者，但这篇文章读起来有点像广告。为了改进答案，您可以透露您隶属于该产品

【参考方案13】：

我曾经为此使用 baron，但现在已切换到 parso，因为它与现代 python 是最新的。它工作得很好。

我也需要这个用于突变测试仪。用 parso 做一个真的很简单，在https://github.com/boxed/mutmut查看我的代码