Easily producing Python AST for multiple expressions in one line

Question

I'd like to recreate this expression with Python AST:

1 == 2; 1 >= 2

And I can achieve this with the following AST structure:

Module(
    body=[
        Expr(value=Compare(left=Num(n=1), ops=[Eq()], comparators=[Num(n=2)])),
        Expr(value=Compare(left=Num(n=1), ops=[GtE()], comparators=[Num(n=2)]))
    ]
)

But the above AST structure is identical for 2 expression in a single line and 2 expressions, each in a separate line.

I know that I can manually calculate and modify the nodes' col_offset and lineno attributes to make it a signle line expression, but is there an easier way?

Answer 1

I assume you've generated the AST by some method other than directly parsing the code. One simple solution might be to

1. convert all or parts of the AST back to a string of python code
2. manipulate that string
3. parse() the result

to get a new effectively equivalent AST with new correct line numbers and column offsets.

For example, walk the full AST, and for every body, convert it to python code, replace all newline characters with semicolons, parse() the result, and replace the original body with the result.

Edit: Here's a simple demo using the astunparse module:

from ast import *
from astunparse import *

a1 = parse("1 == 2\n1 >= 2")
print map(lambda x: (x.lineno, x.col_offset), a1.body)  # [(1, 0), (2, 0)]

a2 = parse(unparse(a1).strip().replace('\n', ';'))
print map(lambda x: (x.lineno, x.col_offset), a2.body)  # [(1, 0), (1, 9)]