Converting string typed Abstract Syntax Tree(AST) to Python Code?

Viewed 491

I am working on a parsed AST string of Python code. Now, I am at a stage where I want to convert small tree structures back to code for some analysis.

import ast
ast_string = ast.dump(ast.parse("[1,2,3]"))
print(ast_string)
# 'Module(body=[Expr(value=List(elts=[Num(n=1), Num(n=2), Num(n=3)], ctx=Load()))])'

Now, I want to convert this Module(body=[Expr(value=List(elts=[Num(n=1), Num(n=2), Num(n=3)], ctx=Load()))]) back to code.

Example:

def ast_to_code(ast_string):
   ....

code = ast_to_code(ast_string)
print(code) # [1,2,3]

Thanks in advance. I have searched over the web found different libraries that take a parse tree object and then convert it back to code. I wasn't able to find something that can convert any AST tree in string back to code.

1 Answers

There are a couple of ways by which the AST string can be converted back to an ast.AST. The first solution is the shortest, since it involves substituting ast.AST subclass names in the original string with their module dot lookup, and then using eval to produce a result:

import ast, re
def ast_to_code(ast_string):
   return ast.unparse(eval(re.sub('\w+(?=\()', lambda x:f'ast.{x.group()}', ast_string)))

ast_string = ast.dump(ast.parse("[1,2,3]"))
print(ast_to_code(ast_string))
#even shorter option with star import:
from ast import *
print(ast.unparse(eval(ast.dump(ast.parse("[1,2,3]")))))

Output:

'[1, 2, 3]'

Of course, the solution above is a bit hackish and rather insecure due to the use of eval. An alternative is to parse the AST string (itself valid Python syntax) to an ast object, then traverse the result and create a new ast object:

import ast
ast_string = ast.dump(ast.parse("[1,2,3]"))
def to_ast(d):
   if isinstance(d, ast.Call):
      return getattr(ast, d.func.id)(*map(to_ast, d.args), 
                     **{i.arg:to_ast(i.value) for i in d.keywords}, **{'lineno': None, 'col_offset': None, 'end_lineno': None, 'end_col_offset': None})
   if isinstance(d, ast.List):
      return list(map(to_ast, d.elts))
   return d.value

def ast_to_code(ast_string):
   return ast.unparse(to_ast(ast.parse(ast_string).body[0].value))

print(ast_to_code(ast_string))

Output:

'[1, 2, 3]'
Related