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?

Are you trying to produce the AST from the source code, or the source code from the AST? What dissatisfies you about the AST you have? — user2357112 supports Monica
Why do you care about how many lines these expressions are spread across, anyway? If you have source code, you can just ast.parse it, and if you don't have source code, you don't really have lines, either, so you might as well fill in dummy values for all the col_offsets and linenos. — user2357112 supports Monica
An AST models the logical content, not the physical line layout. That the AST contains line and column numbers still is a token nod towards rewriting usecases where people want to alter existing code. If you are not generating the AST from existing source, then you are stuck with calculating the line and column info yourself. — Martijn Pieters
I can imagine caring about column offsets and line numbers when creating code rewriting tools that try to maintain the user's formatting as much as possible. It also might matter for ascii art, polyglots, minifying, code golf, etc. I'm curious what @noamt needed it for. — Alex Varga
@AlexVagra I'm taking a shot at writing a testing a Python testing framework that's heavily inspired by Groovy's Spock framework. It's still a WIP - github.com/browncoat-ninjas/nimoy — noamt

Alex Varga Alex Varga · Accepted Answer · 2017-09-07T17:49:47

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

convert all or parts of the AST back to a string of python code
manipulate that string
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)]

Easily producing Python AST for multiple expressions in one line

1 Answers