Convert python objects to python AST-nodes

Question

I have a need to dump the modified python object back into source. So I try to find something to convert real python object to python ast.Node (to use later in astor lib to dump source)

Example of usage I want, Python 2:

import ast
import importlib

import astor


m = importlib.import_module('something')

# modify an object
m.VAR.append(123)

ast_nodes = some_magic(m)

source = astor.dump(ast_nodes)

Please help me to find that some_magic

Answer 1

There's no way to do what you want, because that's not how ASTs work. When the interpreter runs your code, it will generate an AST out of the source files, and interpret that AST to generate python objects. What happen to those objects once they've been generated has nothing to do with the AST.

It is however possible to get the AST of what generated the object in the first place. The module inspect lets you get the source code of some python objects:

import ast
import importlib
import inspect

m = importlib.import_module('pprint')
s = inspect.getsource(m)
a = ast.parse(s)
print(ast.dump(a))
# Prints the AST of the pprint module

But getsource() is aptly named. If I were to change the value of some variable (or any other object) in m, it wouldn't change its source code.

Even if it was possible to regenerate an AST out of an object, there wouldn't be a single solution some_magic() could return. Imagine I have a variable x in some module, that I reassign in another module:

# In some_module.py
x = 0

# In __main__.py
m = importlib.import_module('some_module')
m.x = 1 + 227

Now, the value of m.x is 228, but there's no way to know what kind of expression led to that value (well, without reading the AST of __main__.py but this would quickly get out of hand). Was it a mere literal? The result of a function call?

If you really have to get a new AST after modifying some value of a module, the best solution would be to transform the original AST by yourself. You can find where your identifier got its value, and replace the value of the assignment with whatever you want. For instance, in my small example x = 0 is represented by the following AST:

Assign(targets=[Name(id='x', ctx=Store())], value=Num(n=0))

And to get the AST matching the reassignment I did in __main__.py, I would have to change the value of the above Assign node as the following:

value=BinOp(left=Num(n=1), op=Add(), right=Num(n=227))

If you'd like to go that way, I recommend you check Python's documentation of the AST node transformer (ast.NodeTransformer), as well as this excellent manual that documents all the nodes you can meet in Python ASTs Green Tree Snakes - the missing Python AST docs.

Answer 2

What Vladimir is asking about is certainly useful for compiler optimizations. Indeed, there are ways to accomplish that using the ast library. Here is a simple example demonstrating evaluation of constant functions:

from ast import *
import numpy as np

PURE_FUNS = {'arange' : np.arange}
PROG = '''
A=arange(5)
B=[0, 1, 2, 3, 4]
A[2:3] = 1
C = [A[1], 2, m]
'''

def py_to_ast(o):
    if type(o) == np.ndarray:
        return List(elts=[py_to_ast(e) for e in o], ctx=Load())
    elif type(o) == np.int64:
        return Constant(value=o)
    # Add elifs for more types here
    else:
        assert False

class EvalPureFuns(NodeTransformer):
    def visit_Call(self, node):
        is_const_args = all(type(a) == Constant for a in node.args)
        if node.func.id in PURE_FUNS and is_const_args:
            res = eval(unparse(node), PURE_FUNS)
            return py_to_ast(res)
        return node

node = parse(PROG)    
node = EvalPureFuns().visit(node)    
print(unparse(node))