python - 如何使用 ast.NodeTransformer 将 List Comprehensions 转换为 For Loops？

Question

我正在尝试借助 ast.NodeTransformer 将列表理解转换为 For 循环。之后，我将对 Python 进行类型分析。

我的理解是遇到一个ast节点ast.ListComp，要转成两个ast节点:第一个节点是ast.Assign，第二个是ast.For。但是根据文档 ast.NodeTransformer , ast.ListComp 只能转换为一个 ast节点而不是两个，这意味着我们不能插入ast.Assign和 ast.For一起变成AST。

现在我可以做一个等效的 for 循环。但我还需要一个赋值语句来将列表初始化为 []。我怎样才能避免这个问题？

一个例子可能会有所帮助:

a=[x for x in [1,2,3]] # One node in AST

到

a=[] # One node in AST
for x in [1,2,3]:  # Another one in AST
    a.append(x)

Answer 1

你可以递归遍历AST，当你在赋值语句中遇到列表理解(ast.ListComp)时，你可以移除父赋值对象并插入到body 范围是从理解派生的标准 for 循环:

import ast
def comp_to_expl(tree):
   if hasattr(tree, 'body'):
      i = 0
      while i < len(tree.body):
         if isinstance(a:=tree.body[i], ast.Assign) and isinstance(a.value, ast.ListComp):
            tree.body = tree.body[:i] + \
                [ast.Assign(
                   targets=[ast.Name(id = a.targets[0].id)], value = ast.List(elts = []),
                   lineno = a.lineno
                )] + \
                [ast.For(
                   target = a.value.generators[0].target,
                   iter = a.value.generators[0].iter,
                   body = [ast.Expr(
                             value = ast.Call(
                                       func = ast.Attribute(value = ast.Name(id = a.targets[0].id), attr = 'append', ctx = ast.Load()),
                                       args = [a.value.elt],
                                       keywords = []
                       ))],
                   lineno = a.lineno+1,
                   orelse = [],
                  )] + \
                tree.body[i+1:]   
            i += 1   
         i += 1
   for i in getattr(tree, '_fields', []):
      if isinstance(v:=getattr(tree, i, None), list):
         for i in v: 
            comp_to_expl(i)
      elif isinstance(v, ast.AST):
         comp_to_expl(v)
         
synt = ast.parse('a=[x for x in [1,2,3]]')
comp_to_expl(synt)
print(ast.unparse(synt))

输出:

a = []
for x in [1, 2, 3]:
    a.append(x)