gpt4 book ai didi

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

转载 作者:太空宇宙 更新时间:2023-11-03 11:21:42 24 4
gpt4 key购买 nike

我目前正在研究抽象语法树,使用 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)

这可以用抽象语法树实现吗? (注意:我知道调用的装饰函数在运行代码时会产生相同的结果,但这不是我所追求的;我需要实际输出修改后的代码,并插入比打印语句更复杂的东西).

最佳答案

如果您询问的是如何将节点插入到低级别的 AST 树中,或者更具体地说,是关于如何使用更高级别的工具进行节点插入以遍历 AST 树(例如,a ast.NodeVisitorastor.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)

(请注意,ast.Call 的参数在 Python 3.4 和 3.5+ 之间发生了变化。如果您使用的是旧版本的 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 节点,但我怀疑这会更具挑战性。

关于python - 使用抽象语法树修改 Python 3 代码,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/41764207/

24 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com