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

问题描述：

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

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

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

解决方案 1：

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

这两个工具都使用了lib2to3库，它是 Python 解析器/编译器机制的一种实现，可以在从源代码 -> AST -> 源代码往返时保留源代码中的注释。

如果您想进行更多类似转换的重构，那么rope 项目可能会满足您的需求。

ast模块是您的另一个选择，并且有一个较旧的示例说明如何将语法树“反解析”回代码（使用解析器模块）。但是，在对代码进行 AST 转换并将其转换为代码对象时，该模块更有用。ast

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

解决方案 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

请注意，您可能会丢失确切的格式和注释，因为这些内容不会被保留。

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

解决方案 3：

花了一段时间，但 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 对象。

解决方案 4：

在不同的答案中，我建议使用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 上测试过这一点。

解决方案 5：

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

• 如果您想要生成一个人们会实际使用的 .py 文件，也许这样他们就可以填写表格并获取一个有用的 .py 文件以插入到他们的项目中，那么您就不想将其更改为 AST 并改回，因为您会丢失所有格式（想想通过将相关行组合在一起使 Python 如此易于阅读的空白行）（ast 节点具有lineno和col_offset属性）注释。相反，您可能希望使用模板引擎（例如， Django 模板语言旨在使模板化文本文件变得容易）来自定义 .py 文件，或者使用 Rick Copeland 的MetaPython扩展。

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

• 但在几乎所有情况下，您可能都在尝试做一些动态的事情，而 Python 这样的语言实际上可以很容易地做到这一点，而无需编写新的 .py 文件！如果您扩展您的问题以让我们知道您真正想要完成什么，那么新的 .py 文件可能根本不会涉及答案；我见过数百个 Python 项目在做数百件现实世界的事情，其中​​没有一个需要编写 .py 文件。所以，我必须承认，我有点怀疑您已经找到了第一个好的用例。:-)

更新：既然您已经解释了要做什么，我还是会倾向于在 AST 上进行操作。您需要通过删除整个语句来进行变异，而不是删除文件中的行（这可能会导致半语句因 SyntaxError 而终止），还有什么地方比 AST 更适合这样做呢？

解决方案 6：

借助ast模块，解析和修改代码结构当然是可能的，稍后我将通过示例展示这一点。但是，仅使用模块无法写回修改后的源代码ast。还有其他模块可用于此工作，例如此处的模块。

注意：下面的示例可以作为ast模块使用方面的入门教程，但是有关使用模块的更全面的指南ast可在此处找到，即Green Tree Snakes 教程和模块官方文档ast。

简介ast：

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

只需调用 API 即可解析 Python 代码（以字符串表示）ast.parse()。这将返回抽象语法树 (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 的打印语句转换为 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 ("
USAGE:
     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：

如果您在 2019 年查看此内容，那么您可以使用这个libcs​​t
包。它的语法与 ast 类似。这非常有效，并且保留了代码结构。它对您必须保留注释、空格、换行符等的项目基本上很有帮助。

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

解决方案 8：

我最近创建了一段相当稳定（核心经过了很好的测试）并且可扩展的代码，它可以从ast树中生成代码：https://github.com/paluh/code-formatter。

我正在使用我的项目作为一个小型 vim 插件（我每天都在使用）的基础，所以我的目标是生成真正漂亮且可读的 python 代码。

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

解决方案 9：

另一个答案推荐codegen，它似乎已被 取代。 PyPI 上astor的 版本（撰写本文时为 0.5 版）似乎也有点过时，因此您可以按如下方式安装 的开发版本。astor`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：

我编写了几个实用程序来做这种事情，在每种情况下我选择的工具都是libcst。Instagram 创建这个是为了操作他们的 Python 代码库；例如插入类型注释。诚然，它没有使用 AST，它是一个 CST，但结构非常相似，而且易于使用。

解决方案 11：

不幸的是，上述答案实际上都没有满足这两个条件

• 保留周围源代码的语法完整性（例如，保留注释、其余代码的其他格式）

• 实际使用 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：

我们有类似的需求，但这里的其他答案无法解决这个问题。因此，我们为此创建了一个库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

希望这有帮助！

解决方案 13：

程序转换系统是一种解析源文本、构建 AST 的工具，允许您使用源到源转换（“如果您看到这种模式，就用那种模式替换它”）对其进行修改。此类工具非常适合对现有源代码进行变异，这些变异只是“如果您看到这种模式，就用模式变体替换”。

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

请参阅此SO 答案，了解 DMS 解析的 Python AST 准确捕获注释的示例。DMS 可以更改 AST，并重新生成有效文本，包括注释。您可以要求它使用自己的格式约定（您可以更改这些约定）漂亮地打印 AST，或者执行“保真打印”，使用原始行和列信息最大限度地保留原始布局（插入新代码的布局不可避免地会发生一些变化）。

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

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

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

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

解决方案 14：

我以前用过 baron，但后来改用 parso，因为它与现代 Python 保持同步。不幸的是，Python 对其解析器进行了很大改动，而 parso 到 2024 年还没有跟上。

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

解决方案 15：

最小代码

import ast
code="""
def foo():
   print("hello world")
"""
p=ast.parse(code)
print(ast.unparse(p))