How to use ast.NodeTransformer to transform List Comprehensions to For Loops?

user_6403253 0
Viewed 99

I'm trying to transform List Comprehensions to For Loops with the help of ast.NodeTransformer. After this, I would do a type analysis for Python.

I understand that when an ast node ast.ListComp is encountered, it has to be transformed to two ast nodes: the first node is ast.Assign and the second is ast.For. But according to the document ast.NodeTransformer, ast.ListComp can only be converted into one ast node rather than two, which means we can not insert ast.Assign and ast.For into AST together.

For now I can make an equivalent for loop. But I also need an assignment statement to initialize a list to []. How can I circumvent this problem?

An example may help:

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

to

a=[] # One node in AST
for x in [1,2,3]:  # Another one in AST
    a.append(x)
abstract-syntax-tree list-comprehension python
1 Answers

You can recursively traverse the AST and when you encounter a list comprehension (ast.ListComp) in an assignment statement, you can remove the parent assignment object and insert into the body scope a standard for-loop that is derived from the comprehension:

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))

Output:

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