Better Tree Traversal in Python

Nobody likes a bad Visitor.

Tree traversal is a common but tricky topic in any programming language. Many developers will be familiar with walking document trees for the web, handling nested JSON data, working with parsed grammars, or recursing through good old binary trees. I started developing the code in this article after experimenting with syntax tree analysis in Rust.

If you Google around, the tree traversal code you find on Stack Overflow and elsewhere tends to fall into two camps: the functional, making heavy use of recursion, and the object-oriented, based on the visitor pattern.

We’re going to take a look first at an implementation of the Visitor pattern and describe some of its pros and cons, and then compare it with a functional-flavoured alternative implementation. The last section will describe a hybrid approach that seems to capture the best of both worlds.

I’m going to use Python’s abstract syntax trees (ASTs) by way of data. This avoids the need for external dependencies or complicated set-up. Although it might help, you shouldn’t need to know much about the specifics of ASTs to follow along.

How do you know which road is less travelled by before you’ve travelled it? 🤔

A Visitor Implementation

I’ll see myself out.

The basic visitor pattern comprises a common interface for handling each node of the tree, called into by a separate piece of logic for handling the structure of the tree. This allows new visitors to be created very easily. In the code below, I’ve set up a Visitor abstract base class and a couple of subclasses which can be used to calculate the number of nodes and the maximum depth of the tree.

This should print 12 and 2 for the total nodes and max depth respectively.

There’s nothing very special about this implementation. I’ve deviated slightly from the original specification to permit a node to both enter and leave its node (this makes calculation of the depth very straightforward), and for syntactic convenience I’ve built the walk logic into the base class. I’ve also demonstrated how to build a “Combined” visitor that can walk several visitors at once — this will become more interesting in future sections.

The result:

Total Nodes: 12
Max Depth: 2

The implementation works, but as you use it you quickly run into difficulties efficiently re-using the code. For example, what if I want to know the depth at every node in the tree, rather than just the maximum depth? What about if I want to do the same thing, but annotate each depth value with the name of the node? Or get the total count of all “Name” nodes (variables)? All of these are possible, but difficult to implement without starting a brand new visitor from scratch.

A Functional Implementation

To understand recursion, read this sentence.

Tree structures lend themselves naturally to a functional style thanks to their recursive structure which can be succinctly handled with a recursive function (see tree() below). This functional style tends to be elegant and practical at the expense of readability.

Again, 12 total nodes, and now we can see the total size of each node in the tree, and for fun the length of the source code as well… but where’s the depth calculation?

To explain the code, let’s focus on the main() function.

• To get the total_size function, I’ve used iter_map_agg to produce a function that will walk the tree, producing a constant value of one at at each node representing the node’s size, and then aggregate by summing.
• The name of a node is simply it’s class’ name, and for fun I’ve built a new function called source_length which is the length of the unparsed source code.
• The main_fn is a bit tricky. I’ve set the mapper to be a function which will produce a tuple for each of its input functions (name, source length, and total size). iter_map_agg will then walk the tree, produce this tuple at each node, and collect the result in a final tuple.

See what I mean about the logic being difficult to follow sometimes? However, we’ve got the result we wanted, a snapshot of each node in the tree, as well as the aggregate values we got from our visitor.

Total Nodes: 12
(('Module', 44, 12),
 ('If', 44, 11),
 ('Name', 5, 2),
 ('Load', 0, 1),
 ('If', 26, 8),
 ('Name', 5, 2),
 ('Load', 0, 1),
 ('Expr', 12, 5),
 ('Call', 12, 4),
 ('Name', 4, 2),
 ('Load', 0, 1),
 ('Constant', 6, 1))

Except… what about the depth?

Trying to come up with a function to calculate the depth of a node is quite difficult. I don’t want to re-implement the logic from the tree() iterator to do it, but equally it’s difficult to think of a function that can operate just on a node that will compute the depth — it’s a property of the node’s context, rather than the node itself.

Analogously to the difficulty of code re-use in the visitor pattern, it is difficult (though again, not impossible) using a functional approach to maintain context in the tree — something more naturally solved with state.

A Hybrid Approach

Don’t take this out of context.

So we want some state to track things like depth in the tree, but we want the composability of the functional style. And there was an interesting word there: context. Let’s create a class whose context is managed separately from its calculation. Context can be handled using a context manager, on a method we’ll call enter(), and the result will be returned from a function called visit().

Hooray! Composition and functional flexibility ✔️ A tiny bit of state for our depth calculation ✔️ Best of both worlds?

This time, I’m going to describe the subclasses in turn:

• FnVisitor, as its name would suggest, wraps a single function, allowing it to inherit the enter() method of the base class. We can use it for simple, context-free calculations like the name or source length.
• ConstantVisitor is the simplest of all, and just returns the same value for any node it is given. I’ve used this in the total_size_visitor, mirroring its use in the functional implementation.
• The DepthVisitor is the first interesting one. When the depth visitor is entered, if the node is an If node, it will increment its internal state by one (like in the visitor pattern) then yield to whatever happens next, before decrementing by one on exit. (Note that the visit method has to subtract this value if the node happens to be an If node, because it’s already been incremented).
• The CombinedVisitor acts like the combine function in the functional style, but to work it needs to enter all of its constituent visitors. To make that work, we use another feature from contextlib: the ExitStack, which allows us to “queue” a series of contexts.
• The TreeVisitor behaves a bit like the iter_map_agg function of the functional approach. By default, it will return a tuple of the result of its sub-visitor applied to all sub-nodes, but we can change the aggregation function, for example to the max or sum of the values, as shown in the main script. When we enter it, all we need to do is enter its sub-visitor.

The result gives us everything we want (assuming your wants are limited to flexible tree traversal): the totals are as easy to calculate as the subtotals, composability, state-friendly calculations like depth, and a clear interface to extend. The result?

Total Nodes: 12
Max Depth: 2
(('Module', 44, 12, 0),  # source length, size, depth
 ('If', 44, 11, 0),
 ('Name', 5, 2, 1),
 ('Load', 0, 1, 1),
 ('If', 26, 8, 1),
 ('Name', 5, 2, 2),
 ('Load', 0, 1, 2),
 ('Expr', 12, 5, 2),
 ('Call', 12, 4, 2),
 ('Name', 4, 2, 2),
 ('Load', 0, 1, 2),
 ('Constant', 6, 1, 2))

You can easily play around with the aggregation functions and the CombinedVisitor to get some other useful results, like dictionaries, annotated calculations (CombinedVisitor(ConstantVisitor(“depth"), DepthVisitor())) and so on.

Where To Next?

• I’ve added a scattering of type hints through the code here, but extending the type annotations with Generics would provide a powerful extension to the composability of these classes (although it might mean a lot of overloading).
• What if Visitor were callable instead of implementing the visit method? It would integrate the functional and stateful styles even further, with functions being valid visitors if they didn’t need any internal state.
• How would you adapt the code for other types of trees? Could we implement a generic set of visitors that could (maybe dynamically) handle any tree structure?
• How would these concepts be implemented in other programming languages? I often find this is a useful exercise for generating new structures or patterns. In Rust, for example, you have to be careful about the type and lifetime of nodes as you recurse a method.
• Can you think of anything else?

Conclusion

The visitor pattern is great for extensibility but difficult to justify in terms of reusability. A functional style improves reusability, but makes certain types of calculation prohibitively obtuse. I’ve demonstrated a hybrid style which has aspects of both styles. Which is your favourite? Do you favour pragmatism or elegance in your code? Are there clever tricks in the functional or visitor style I could have used to solve the problems?

Thank you for reading!