perimosocordiae perimosocordiae - 6 months ago 13
Python Question

Idiomatic Python: Propagating yields or flattening sequences?

I'm writing a breadth depth-first tree traversal function, and what I want to do is this:

def traverse(node):
yield node
for n in node.children:
yield_all traverse(n) # << if Python had a yield_all statement


The idea is to end up with a (flat) sequence of nodes in the tree.

Approach #1: (propagating yields)

def traverse(node):
yield node
for n in node.children:
for m in traverse(n):
yield m


Approach #2: (flattening sequences)

def traverse(node):
return itertools.chain([node],*(traverse(n) for n in node.children))


The first approach seems more clean, but I feel weird explicitly
yield
ing each node in the subtree at each level.

The second approach is terse and slightly dirty, but it matches what I would write in Haskell:

traverse node = node : concatMap traverse (children node)


So my question is: Which is better? Or am I missing a best 3rd option?

Answer

You are not the only one that pines for a yield_all, see PEP380 (draft status for nearly 2 years, sigh).

This won't get me any Pythonic medal, but since you are clearly a functional-programming-guy, you could always write this (and feel at home):

import itertools

def concatMap(fun, it):
    return itertools.chain.from_iterable(itertools.imap(fun, it))

def traverse(node):
    return itertools.chain([node], concatMap(traverse, node.children))

Incidentally, I wrote a short Python implementation for PEP380 some time ago, using it it would look like this:

from pep380 import supergenerator, _from

@supergenerator
def traverse(node):
    yield node
    for n in node.children:
      yield _from(traverse(n))

Of course, this is just a curiosity, if I expected someone else to read the code, I'd probably use your first approach (verbose as it is).

Comments