CODEWITHSUNDEEP
pythonpython-3.xabstract-syntax-tree
samuelbrody1249
samuelbrody1249

Reputation: 4767

How to format an AST parse

I would like to format the following ast parse:

>>> import ast
>>> print(ast.dump(ast.parse('-a+b')))
Module(body=[Expr(value=BinOp(left=UnaryOp(op=USub(), operand=Name(id='a', ctx=Load())), op=Add(), right=Name(id='b', ctx=Load())))])

It seems like the indent option was introduced in python3.9, but I don't see an option to 'pretty-print' before then. What options are there to print a nicely-formatting output for an the syntax tree?

Upvotes: 2

Views: 2398

Answers (2)

Paulo Almeida
Paulo Almeida

Reputation: 2154

I had one use case in which I couldn't upgrade to Python 3.9 (where indent argument was added), yet I needed a way to prettify the result of ast.dump.

I wrote the following method that can take an unformatted ast.dump output and print that in a way that is easier on the eyes.

def prettify(ast_tree_str, indent=4):
    ret = []
    stack = []
    in_string = False
    curr_indent = 0

    for i in range(len(ast_tree_str)):
        char = ast_tree_str[i]
        if in_string and char != '\'' and char != '"':
            ret.append(char)
        elif char == '(' or char == '[':
            ret.append(char)

            if i < len(ast_tree_str) - 1:
                next_char = ast_tree_str[i+1]
                if next_char == ')' or next_char == ']':
                    curr_indent += indent
                    stack.append(char)
                    continue

            print(''.join(ret))
            ret.clear()

            curr_indent += indent
            ret.append(' ' * curr_indent)
            stack.append(char)
        elif char == ',':
            ret.append(char)

            print(''.join(ret))
            ret.clear()
            ret.append(' ' * curr_indent)
        elif char == ')' or char == ']':
            ret.append(char)
            curr_indent -= indent
            stack.pop()
        elif char == '\'' or char == '"':

            if (len(ret) > 0 and ret[-1] == '\\') or (in_string and stack[-1] != char):
                ret.append(char)
                continue

            if len(stack) > 0 and stack[-1] == char:
                ret.append(char)
                in_string = False
                stack.pop()
                continue

            in_string = True
            ret.append(char)
            stack.append(char)
        elif char == ' ':
            pass
        else:
            ret.append(char)

    print(''.join(ret))

Usage:

if __name__ == '__main__':
    content = """
@testdecorator
def my_method(a, b):
    def ola():
        print("Hello")
    ola()
    return (a + b) * 5 + "dasdas,da'sda\\'asdas\\'\\'"
    """

    ast_tree = ast.parse(source=content)
    prettify(ast.dump(ast_tree))

PS.: It's not 100% equivalent to what one could get out of the Python 3.9 ast.dump(...., indent=n) but it should be enough for now. Feel free to improve it

Upvotes: 2

lvrf
lvrf

Reputation: 308

If you need to pretty-print the AST in an earlier python version and are happy with the indent function in Python3.9 why not just take the dump function from 3.9 and implement it in your project? The source code is here: https://github.com/python/cpython/blob/e56d54e447694c6ced2093d2273c3e3d60b36b6f/Lib/ast.py#L111-L175

And it doesn't look very complicated and doesn't seem to use any features specific to 3.9.

Upvotes: 4

Related Questions