December 8th 2006

Visitor pattern in Python

One design pattern I'm using a lot in my compiler is the Visitor pattern.

This pattern allows different operations to be performed on the elements in a hierarchical object structure. Crucially though, it allows you to define a new set of operations without changing the classes of the elements on which it operates and thus provides a nice way of seperating your data and algorithms.

So, how would we use it? Let's imagine we want to traverse an Abstract Syntax Tree for a simple language, first to generate code for some stack-based machine and secondly to make an image of the tree. We could add suitable code to our existing ASTNode subclasses, but this would quickly get messy. In idiomatic Java, we might setup the interface and a concrete implementation like this:

interface IVisitor {
    void accept(ASTNode n);
    void accept(EchoStatement n);
    void accept(BinaryExpression n);
    void accept(Constant n);
}
class CodeGeneratorVisitor implements IVisitor {
    public void accept(ASTNode n) {
        // Traverse all child nodes
        for (ASTNode child : n.children) {
            self.accept(child);
        }
    }
    public void accept(BinaryExpression n) {
        // ... etc
    }
}

To stop it getting too verbose, we might then refactor out the 'for-each' which loop around the child nodes to a base class for all our concrete IVisitor implementations, or possibly convert IVisitor to an abstract class and mark the current methods as abstract. (This would also give us the opportunity to specify a default implementation for an unhandled node, but veers far from the classic design)

In Python however, we obviously can't provide a solution like this as we are dealing with a different typing model. One solution that is often suggested is to add a method to each of the elements to be traversed. This is passed the concrete visitor when being called and they simply call type-specific methods back in the visitor. Although this would work, it doesn't sound like seperating the data and the algorithm to me.

A much more Python-like solution would be to use an "Extrinsic Visitor". Here, type introspection is used to call an appropriate method based on it's name: for example, visiting a node of type Spam would call visit_Spam. This works fairly well until you want to provide the same functionality to all subclasses of a class: it would always call the function with the name of subclass. Some hacks could of course be added to allow a fallback to a parent class, but it wouldn't exactly be ideal.

Step forward generic functions, a great idea taken from Lisp. In our Python implementation, calls to the generic function will result in program execution being dispatched to the function that is decorated with a compatible type of the argument.

Lets first setup our AST and then we can see this single-dispatch model in action. Here is the original code:

echo 1;
echo 2 + (4 * 3);

and is represented in our AST thus:

 >>> tree = ASTNode([
 EchoStatement(
     Constant([], {'value': 1})
 ),
 EchoStatement(
     BinaryExpression([
         Constant([], {'value': 2}),
         BinaryExpression([
             Constant([], {'value': 4}),
             Constant([], {'value': 3})
         ], {'operator': 'multiply'})
     ], {'operator': 'plus'})
  )
])

EchoStatement, Constant and BinaryExpression are subclasses of ASTNode which store an ordered list of child nodes in self.children and dictionary of properties in self.props. We can now generate code for a stack machine like this:

>>> class CodeGeneratorVisitor(object):
    @dispatch.on('node')
    def visit(self, node):
        """This is the generic method"""

    @visit.when(ASTNode)
    def visit(self, node):
        map(self.visit, node.children)

    @visit.when(EchoStatement)
    def visit(self, node):
        self.visit(node.children)
        print "print"

    @visit.when(BinaryExpression)
    def visit(self, node):
        map(self.visit, node.children)
        print node.props['operator']

    @visit.when(Constant)
    def visit(self, node):
        print "push %d" % node.props['value']

>>> CodeGeneratingVisitor().visit(tree)
push 1
print
push 2
push 4
push 3
multiply
plus
print

Bog-standard post-order traversal but note the brevity. This way is not only cleaner and gives us the benefit of sub-class compatibile type checking, it's at least 2 or 3 times faster as it avoids string concatenation, introspection and the dispatching code runs inside C.

Graphing is slightly more involved but still succinct:

class GraphingVisitor(object):
    def labelNode(self, node, str):
        print '"%s" [ label="%s" ];' % (id(node), str)

    @dispatch.on('node')
    def visit(self, node): pass

    @visit.when(ASTNode)
    def visit(self, node):
        print 'digraph "g" {'
        for child in node.children:
            print 'Root ->',
            self.visit(child)
        print '}'

    @visit.when(EchoStatement)
    def visit(self, node):
        print '"%s";' % id(node)
        self.labelNode(node, 'echo')
        print '"%s" -> ' % id(node),
        self.visit(node.children)

    @visit.when(BinaryExpression)
    def visit(self, node):
        print '"%s";' % id(node)
        self.labelNode(node, node.props['operator'])
        for child in node.children:
            print '"%s" -> ' % id(node),
            self.visit(child)

    @visit.when(Constant)
    def visit(self, node):
        print '"%s";' % id(node)
        self.labelNode(node, node.props['value'])

>>> GraphingVisitor().visit(tree)
Root -> "-1480985620";
"-1480985620" [ label="echo" ];
"-1480985620" ->  "-1480985588";
"-1480985588" [ label="1" ];
Root -> "-1480985140";
"-1480985140" [ label="echo" ];
"-1480985140" ->  "-1480985172";
"-1480985172" [ label="plus" ];
"-1480985172" ->  "-1480985300";
"-1480985300" [ label="2" ];
"-1480985172" ->  "-1480985204";
"-1480985204" [ label="multiply" ];
"-1480985204" ->  "-1480985268";
"-1480985268" [ label="4" ];
"-1480985204" ->  "-1480985236";
"-1480985236" [ label="3" ];

Which, when run though dot gives this shiney graph.

Anyway, I hope this gives a little bit of insight into the Visitor pattern in Python. I've simplified a few things somewhat, but I'm essentially using the same code in my project with a number of concrete visitors performing tasks such as type checking, constant folding translating various syntactic sugars into their semantic equivalents, removing dead code, etc.




You can subscribe to new posts via email or RSS.