biais.org

Python standard library allows developper to manipulate Python Abstract Syntax Tree (AST) and bytecode objects. I wondered how to write a Python code pretty printer, I think the best way is to use the ASTVisitor from compiler.visitor module. Warning, the following code is only for demonstration purposes; for example visitConst will badly process string constants:

import compiler
 
class CodePrinter:
    def __init__(self):
        self.src = ''
 
    def visitName(self,t):
        self.src += t.name
 
    def visitConst(self,t):
        self.src += str(t.value)
 
    def visitStmt(self, t):
        for i in t:
            a = pretty_print(i)
            self.src += a + "\n"
 
    def visitAssName(self, t):
        self.src += t.name + " = "
 
def pretty_print(node):
    myvisitor = CodePrinter()
    # compiler.walk return the visitor instance : 2nd arg
    return compiler.walk(node, myvisitor).src

A simple example to use the previous pretty printer

test = """hello =              1; hey = 2
ola = 3"""
# Parse the test string, return the AST of the parsed string
node = compiler.parse(test)
print pretty_print(node)

produces the following output:

hello = 1
hey = 2
ola = 3

All statements (visitStmt) now finish by a newline and successive white spaces disapeared.

Now, you want to write your own visitor but you don't want to miss one of the AST node. You may try to use the walker class ExampleASTVisitor from compiler.visitor instead of the default one (ASTVisitor). In verbose mode, ExampleASTVisitor will produce a text output of node you may forget to visit. Let modify the pretty_print function:

def pretty_print(node):
    myvisitor = CodePrinter()
    example_walker = compiler.visitor.ExampleASTVisitor()
    example_walker.VERBOSE = 1
    return compiler.walk(node, myvisitor, walker=example_walker).src

The previous example with the modified function produces:

<__main__.CodePrinter instance at 0xb7cc306c>
compiler.ast.Module
     [...]
<__main__.CodePrinter instance at 0xb7cc31ac>
compiler.ast.Assign
     [...]

hello = 1
hey = 2
ola = 3

It seems we forgot to visit Module and Assign nodes. Let's print the AST of the parsed code:

>>> test = """hello =              1; hey = 2
... ola = 3"""
>>> compiler.parse(test)
Module(None, 
  Stmt([
    Assign([AssName('hello', 'OP_ASSIGN')], Const(1)), 
    Assign([AssName('hey', 'OP_ASSIGN')], Const(2)), 
    Assign([AssName('ola', 'OP_ASSIGN')], Const(3))
  ])
)

In this print we see 1 Module and 3 Assign nodes. If we need them in our visitor, we have to had visitModule(self, t) and visitAssign(self, t) methods.