Search code examples
pythonabstract-syntax-tree

How to get source corresponding to a Python AST node?


Python AST nodes have lineno and col_offset attributes, which indicate the beginning of respective code range. Is there an easy way to get also the end of the code range? A 3rd party library?


Solution

  • We had a similar need, and I created the asttokens library for this purpose. It maintains the source in both text and tokenized form, and marks AST nodes with token information, from which text is also readily available.

    It works with Python 2 and 3 (tested with 2.7 and 3.5). For example:

    import ast, asttokens
    st='''
    def greet(a):
      say("hello") if a else say("bye")
    '''
    atok = asttokens.ASTTokens(st, parse=True)
    for node in ast.walk(atok.tree):
      if hasattr(node, 'lineno'):
        print atok.get_text_range(node), node.__class__.__name__, atok.get_text(node)
    

    Prints

    (1, 50) FunctionDef def greet(a):
      say("hello") if a else say("bye")
    (17, 50) Expr say("hello") if a else say("bye")
    (11, 12) Name a
    (17, 50) IfExp say("hello") if a else say("bye")
    (33, 34) Name a
    (17, 29) Call say("hello")
    (40, 50) Call say("bye")
    (17, 20) Name say
    (21, 28) Str "hello"
    (40, 43) Name say
    (44, 49) Str "bye"