Discussion forum for David Beazley

@singledispatch and the visitor pattern

#1

When writing a compiler, it’s common to build various sorts of abstract syntax trees. To do this, you’ll often define various classes like this:

class Node:
    pass

class UnaryOperator(Node):
    def __init__(self, operand):
        self.operand = operand

class BinaryOperator(Node):
    def __init__(self, left, right):
        self.left = left
        self.right = right

class Add(BinaryOperator):
    pass

class Sub(BinaryOperator):
    pass

class Mul(BinaryOperator):
    pass

class Div(BinaryOperator):
    pass

class Negate(UnaryOperator):
    pass

class Number(Node):
    def __init__(self, value):
        self.value = value

Here’s an example of an expression encoded with these classes:

# Representation of 1 + 2 * (3 - 4) / 5                                                                                  
t1 = Sub(Number(3), Number(4))
t2 = Mul(Number(2), t1)
t3 = Div(t2, Number(5))
t4 = Add(Number(1), t3)

To process a syntax tree, you then use a variant of the so-called “visitor” pattern. For example, this code illustrates a common recipe for it (this is taken almost directly from the Python Cookbook).

class NodeVisitor:
    def visit(self, node):
        methname = 'visit_' + type(node).__name__
        meth = getattr(self, methname, None)
        if meth is None:
            meth = self.generic_visit
        return meth(node)

    def generic_visit(self, node):
        raise RuntimeError('No {} method'.format('visit_' + type(node).__name__))

class StackCode(NodeVisitor):
    def generate_code(self, node):
        self.instructions = []
        self.visit(node)
        return self.instructions

    def visit_Number(self, node):
        self.instructions.append(('PUSH', node.value))

    def visit_Add(self, node):
        self.visit(node.left)
        self.visit(node.right)
        self.instructions.append(('ADD',))

    def visit_Sub(self, node):
        self.visit(node.left)
        self.visit(node.right)
        self.instructions.append(('SUB',))

    def visit_Mul(self, node):
        self.visit(node.left)
        self.visit(node.right)
        self.instructions.append(('MUL',))

    def visit_Div(self, node):
        self.visit(node.left)
        self.visit(node.right)
        self.instructions.append(('DIV',))

    def visit_Negate(self, node):
        self.visit(node.operand)
        self.instructions.append(('NEG',))

# Generate instructions for a stack machine
s = StackCode()
code = s.generate_code(t4)
for c in code:
    print(c)

This technique for processing a tree is the same used as Python’s own ast module. It’s also similar to something you might find in a Design Patterns book.

Lately, I’ve been experimenting with using the @singledispatch decorator. It turns out that you can use that to greatly simplify the code. Here is a version that uses nothing more than functions and a few other tricks.

from functools import singledispatch

def stack_code(node):
    instructions = [ ]

    @singledispatch
    def visit(node):
        raise RuntimeError("No matching visit() function")

    @visit.register
    def _(node: Number):
        instructions.append(('PUSH', node.value))

    @visit.register
    def _(node:BinaryOperator):
        visit(node.left)
        visit(node.right)
        instructions.append((type(node).__name__.upper(),))

    @visit.register
    def _(node:Negate):
        visit(node.operand)
        instructions.append(('NEG',))

    visit(node)
    return instructions

# --- Example of using the function

code = stack_code(t4)
for c in code:
    print(c)

As you can see, the code is much smaller and there is much less plumbing. There are some interesting bits with the type-matching. For example, you can write a single function to match an abstract class such as BinaryOperator and have it match any of its subclasses (e.g., Add, Mul, etc.). This is not possible in the formulation of the classic visitor class. One limitation of @singledispatch is that it can’t be used with methods of a class. However, if you’re willing to use nested functions instead, you can probably work around it.

On the whole, I’m sold. @singledispatch for the win!

5 Likes