Parse a .py file, read the AST, modify it, then wr

Question

I want to programmatically edit python source code. Basically I want to read a .py file, generate the AST, and then write back the modified python source code (i.e. another .py file).

There are ways to parse/compile python source code using standard python modules, such as ast or compiler. However, I don't think any of them support ways to modify the source code (e.g. delete this function declaration) and then write back the modifying python source code.

UPDATE: The reason I want to do this is I'd like to write a Mutation testing library for python, mostly by deleting statements / expressions, rerunning tests and seeing what breaks.

Answer 1

I used to use baron for this, but have now switched to parso because it's up to date with modern python. It works great.

I also needed this for a mutation tester. It's really quite simple to make one with parso, check out my code at https://github.com/boxed/mutmut

Answer 2

One of the other answers recommends codegen, which seems to have been superceded by astor. The version of astor on PyPI (version 0.5 as of this writing) seems to be a little outdated as well, so you can install the development version of astor as follows.

pip install git+https://github.com/berkerpeksag/astor.git#egg=astor

Then you can use astor.to_source to convert a Python AST to human-readable Python source code:

>>> import ast
>>> import astor
>>> print(astor.to_source(ast.parse('def foo(x): return 2 * x')))
def foo(x):
    return 2 * x

I have tested this on Python 3.5.

Answer 3

We had a similar need, which wasn't solved by other answers here. So we created a library for this, ASTTokens, which takes an AST tree produced with the ast or astroid modules, and marks it with the ranges of text in the original source code.

It doesn't do modifications of code directly, but that's not hard to add on top, since it does tell you the range of text you need to modify.

For example, this wraps a function call in WRAP(...), preserving comments and everything else:

example = """
def foo(): # Test
  '''My func'''
  log("hello world")  # Print
"""

import ast, asttokens
atok = asttokens.ASTTokens(example, parse=True)

call = next(n for n in ast.walk(atok.tree) if isinstance(n, ast.Call))
start, end = atok.get_text_range(call)
print(atok.text[:start] + ('WRAP(%s)' % atok.text[start:end])  + atok.text[end:])

Produces:

def foo(): # Test
  '''My func'''
  WRAP(log("hello world"))  # Print

Hope this helps!

Answer 4

In a different answer I suggested using the astor package, but I have since found a more up-to-date AST un-parsing package called astunparse:

>>> import ast
>>> import astunparse
>>> print(astunparse.unparse(ast.parse('def foo(x): return 2 * x')))


def foo(x):
    return (2 * x)

I have tested this on Python 3.5.

Answer 5

Parsing and modifying the code structure is certainly possible with the help of ast module and I will show it in an example in a moment. However, writing back the modified source code is not possible with ast module alone. There are other modules available for this job such as one here.

NOTE: Example below can be treated as an introductory tutorial on the usage of ast module but a more comprehensive guide on using ast module is available here at Green Tree snakes tutorial and official documentation on ast module.

Introduction to ast:

>>> import ast
>>> tree = ast.parse("print 'Hello Python!!'")
>>> exec(compile(tree, filename="<ast>", mode="exec"))
Hello Python!!

You can parse the python code (represented in string) by simply calling the API ast.parse(). This returns the handle to Abstract Syntax Tree (AST) structure. Interestingly you can compile back this structure and execute it as shown above.

Another very useful API is ast.dump() which dumps the whole AST in a string form. It can be used to inspect the tree structure and is very helpful in debugging. For example,

On Python 2.7:

>>> import ast
>>> tree = ast.parse("print 'Hello Python!!'")
>>> ast.dump(tree)
"Module(body=[Print(dest=None, values=[Str(s='Hello Python!!')], nl=True)])"

On Python 3.5:

>>> import ast
>>> tree = ast.parse("print ('Hello Python!!')")
>>> ast.dump(tree)
"Module(body=[Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Str(s='Hello Python!!')], keywords=[]))])"

Notice the difference in syntax for print statement in Python 2.7 vs. Python 3.5 and the difference in type of AST node in respective trees.

---

How to modify code using ast:

Now, let's a have a look at an example of modification of python code by ast module. The main tool for modifying AST structure is ast.NodeTransformer class. Whenever one needs to modify the AST, he/she needs to subclass from it and write Node Transformation(s) accordingly.

For our example, let's try to write a simple utility which transforms the Python 2 , print statements to Python 3 function calls.

Print statement to Fun call converter utility: print2to3.py:

#!/usr/bin/env python
'''
This utility converts the python (2.7) statements to Python 3 alike function calls before running the code.

USAGE:
     python print2to3.py <filename>
'''
import ast
import sys

class P2to3(ast.NodeTransformer):
    def visit_Print(self, node):
        new_node = ast.Expr(value=ast.Call(func=ast.Name(id='print', ctx=ast.Load()),
            args=node.values,
            keywords=[], starargs=None, kwargs=None))
        ast.copy_location(new_node, node)
        return new_node

def main(filename=None):
    if not filename:
        return

    with open(filename, 'r') as fp:
        data = fp.readlines()
    data = ''.join(data)
    tree = ast.parse(data)

    print "Converting python 2 print statements to Python 3 function calls"
    print "-" * 35
    P2to3().visit(tree)
    ast.fix_missing_locations(tree)
    # print ast.dump(tree)

    exec(compile(tree, filename="p23", mode="exec"))

if __name__ == '__main__':
    if len(sys.argv) <=1:
        print ("\nUSAGE:\n\t print2to3.py <filename>")
        sys.exit(1)
    else:
        main(sys.argv[1])

This utility can be tried on small example file, such as one below, and it should work fine.

Test Input file : py2.py

class A(object):
    def __init__(self):
        pass

def good():
    print "I am good"

main = good

if __name__ == '__main__':
    print "I am in main"
    main()

Please note that above transformation is only for ast tutorial purpose and in real case scenario one will have to look at all different scenarios such as print " x is %s" % ("Hello Python").

Answer 6

Pythoscope does this to the test cases it automatically generates as does the 2to3 tool for python 2.6 (it converts python 2.x source into python 3.x source).

Both these tools uses the lib2to3 library which is a implementation of the python parser/compiler machinery that can preserve comments in source when it's round tripped from source -> AST -> source.

The rope project may meet your needs if you want to do more refactoring like transforms.

The ast module is your other option, and there's an older example of how to "unparse" syntax trees back into code (using the parser module). But the ast module is more useful when doing an AST transform on code that is then transformed into a code object.

The redbaron project also may be a good fit (ht Xavier Combelle)