Search code examples
python-3.xabstract-syntax-treepython-ast

ast nodes not preserving some properties (lineno or/and col_offset)


I'm trying to convert every break statement with exec('break') in a code. So far I've got this:

import ast

source = '''some_list = [2, 3, 4, 5]
for i in some_list:
    if i == 4:
        p = 0
        break
exec('d = 9')'''
tree = ast.parse(source)

class NodeTransformer(ast.NodeTransformer):
    def visit_Break(self, node: ast.Break):
        print(ast.dump(node))
        exec_break = ast.Call(func=ast.Name(id='exec', ctx=ast.Load()),
                              args=[ast.Constant(value='break')],
                              keywords=[])
        return ast.copy_location(exec_break, node)

NodeTransformer().visit(tree)
print(ast.unparse(tree))

However, at the end it outputs p = 0 and exec('break') at the same line:

some_list = [2, 3, 4, 5]
for i in some_list:
    if i == 4:
        p = 0exec('break')
exec('d = 9')

I created the ast.Call object to the exec function with first argument 'break' but it seems not to transform properly. What did I miss?


Solution

  • I've found the bug. The ast.Call node has to be an ast.Expr object:

    def visit_Break(self, node: ast.Break):
        exec_break = ast.Call(func=ast.Name(id='exec', ctx=ast.Load()),
                              args=[ast.Constant(value='break')],
                              keywords=[])
        new_node = ast.Expr(value=exec_break)
        ast.copy_location(new_node, node)
        ast.fix_missing_locations(new_node)
        return new_node
    

    Reference: https://greentreesnakes.readthedocs.io/en/latest/examples.html#simple-test-framework