使用抽象语法树修改Python 3代码

Question

我目前正在使用 ast 和 astor 模块来研究抽象语法树。该文档教我如何检索和漂亮打印各种函数的源代码，网络上的各种示例展示了如何通过将一行内容替换为另一行内容或将所有出现的 + 更改为 * 来修改部分代码。

但是，我想在不同的地方插入额外的代码，特别是当一个函数调用另一个函数时。例如，以下假设函数：

def some_function(param):
    if param == 0:
       return case_0(param)
    elif param < 0:
       return negative_case(param)
    return all_other_cases(param)

会产生（一旦我们使用过astor.to_source(modified_ast)）：

def some_function(param):
    if param == 0:
       print ("Hey, we're calling case_0")
       return case_0(param)
    elif param < 0:
       print ("Hey, we're calling negative_case")
       return negative_case(param)
    print ("Seems we're in the general case, calling all_other_cases")
    return all_other_cases(param)

这对于抽象语法树来说可能吗？（注意：我知道运行代码时调用的装饰函数会产生相同的结果，但这不是我想要的；我需要实际输出修改后的代码，并插入比 print 语句更复杂的东西）。

Answer 1

从您的问题中不清楚您是否询问如何在低级别将节点插入到 AST 树中，或者更具体地说是如何使用更高级别的工具进行节点插入来遍历 AST 树（例如 或 的子类ast.NodeVisitor）astor.TreeWalk）。

在低层插入节点非常容易。您只需list.insert在树中的适当列表上使用即可。例如，这里有一些代码添加了print您想要的三个调用中的最后一个（其他两个几乎同样简单，它们只是需要更多索引）。大部分代码正在为打印调用构建新的 AST 节点。实际的插入非常短：

source = """
def some_function(param):
    if param == 0:
       return case_0(param)
    elif param < 0:
       return negative_case(param)
    return all_other_cases(param)
"""

tree = ast.parse(source) # parse an ast tree from the source code

# build a new tree of AST nodes to insert into the main tree
message = ast.Str("Seems we're in the general case, calling all_other_cases")
print_func = ast.Name("print", ast.Load())
print_call = ast.Call(print_func, [message], []) # add two None args in Python<=3.4
print_statement = ast.Expr(print_call)

tree.body[0].body.insert(1, print_statement) # doing the actual insert here!

# now, do whatever you want with the modified ast tree.
print(astor.to_source(tree))

输出将是：

def some_function(param):
    if param == 0:
        return case_0(param)
    elif param < 0:
        return negative_case(param)
    print("Seems we're in the general case, calling all_other_cases")
    return all_other_cases(param)

（请注意，Python 3.4 和 3.5+ 之间的参数发生了ast.Call变化。如果您使用的是旧版本的 Python，则可能需要添加两个额外的None参数ast.Call(print_func, [message], [], None, None)：）

如果您使用更高级别的方法，事情会有点棘手，因为代码需要弄清楚在哪里插入新节点，而不是使用您自己的输入知识来硬编码事物。

这是一个子类的快速而肮脏的实现 ，它将打印调用添加为在其下TreeWalk有节点的任何语句之前的语句。Call请注意，Call节点包括对类的调用（以创建实例），而不仅仅是函数调用。此代码仅处理一组嵌套调用的最外层，因此如果代码已foo(bar())插入，print则只会提及foo：

class PrintBeforeCall(astor.TreeWalk):
    def pre_body_name(self):
        body = self.cur_node
        print_func = ast.Name("print", ast.Load())
        for i, child in enumerate(body[:]):
            self.__name = None
            self.walk(child)
            if self.__name is not None:
                message = ast.Str("Calling {}".format(self.__name))
                print_statement = ast.Expr(ast.Call(print_func, [message], []))
                body.insert(i, print_statement)
        self.__name = None
        return True

    def pre_Call(self):
        self.__name = self.cur_node.func.id
        return True

你可以这样称呼它：

source = """
def some_function(param):
    if param == 0:
       return case_0(param)
    elif param < 0:
       return negative_case(param)
    return all_other_cases(param)
"""

tree = ast.parse(source)

walker = PrintBeforeCall()   # create an instance of the TreeWalk subclass
walker.walk(tree)   # modify the tree in place

print(astor.to_source(tree)

这次的输出是：

def some_function(param):
    if param == 0:
        print('Calling case_0')
        return case_0(param)
    elif param < 0:
        print('Calling negative_case')
        return negative_case(param)
    print('Calling all_other_cases')
    return all_other_cases(param)

这并不完全是您想要的消息，但也很接近了。walker 无法详细描述正在处理的情况，因为它只查看正在调用的名称函数，而不是到达那里的条件。如果您有一组定义明确的事物要查找，您也许可以更改它以查看节点ast.If，但我怀疑这会更具挑战性。