CODEWITHSUNDEEP
pythonabstract-syntax-tree
Marc Fawver
Marc Fawver

Reputation: 13

Function isn't added to new line, when adding node to AST in Python

in Python I try to add a print statement after each for-loop in a source code using the AST. The problem is however, that the print statement isn't added to a new line but is added at the same line, as the for-loop. Adding various combinations of fix_missing_locations() and increment_lineno() didn't help. What am I doing wrong?

import astor
import ast

class CodeInstrumentator(ast.NodeTransformer):
    def get_print_stmt(self, lineno):
        return ast.Call(
            func=ast.Name(id='print', ctx=ast.Load()),
            args=[ast.Num(n=lineno)],
            keywords=[]
            )

    def insert_print(self, node):
        node.body.insert(0, self.get_print_stmt(node.lineno))

    def visit_For(self, node):
        self.insert_print(node)
        self.generic_visit(node)
        return node

def main():
    input_file = 'source.py'
    try:
        myAST = astor.parsefile(input_file)
    except Exception as e:
        raise e

    CodeInstrumentator().visit(myAST)
    instru_source = astor.to_source(myAST)
    source_file = open('test.py', 'w')
    source_file.write(instru_source)

if __name__ == "__main__":
    main()

Upvotes: 1

Views: 1318

Answers (1)

Cristian Ramon-Cortes
Cristian Ramon-Cortes

Reputation: 1888

This question seems abandoned by I was facing a similar problem and I finally find out a solution so I am writing it down just in case it is useful for somebody.

First of all, notice that ASTOR does not rely nor on lineno nor col_offset so using ast.fix_missing_locations(node), increment_lineno(node, n=1) or new_node = ast.copy_location(new_node, node) will not have any effect on the output code.

This said, the problem is that the Call statement is not a standalone operation and, thus, ASTOR applies it to the previous node (as it was part of the same operation but you have miss-written the node's lineno).

Then, the solution is to wrap the Call statement with a void call using the Expr statement:

def get_print_stmt(self, lineno):
    return ast.Expr(value=ast.Call(
        func=ast.Name(id='print', ctx=ast.Load()),
        args=[ast.Num(n=lineno)],
        keywords=[]
        ))

If you write a code containing a void call to a function you will notice that its AST representation already contains the Expr node:

test_file.py

#!/usr/bin/python

# -*- coding: utf-8 -*-

#
# MAIN
#

my_func()

process_file.py

#!/usr/bin/python

# -*- coding: utf-8 -*-

from __future__ import print_function

def main():
    tree = astor.code_to_ast.parse_file("test_file.py")

    print("DUMP TREE")
    print(astor.dump_tree(tree))
    print("SOURCE")
    print(astor.to_source(tree))

#
# MAIN
#

if __name__ == '__main__':
    main()

OUTPUT

$ python process_file.py

DUMP TREE
Module(
    body=[
        Expr(value=Call(func=Name(id='my_func'), args=[], keywords=[], starargs=None, kwargs=None))])
SOURCE
my_func()

Upvotes: 2

Related Questions