Python Cookbook

Flattening a Nested Sequence

Credit: Luther Blissett

Problem

You have a sequence, such as a list, some of whose items may in turn be lists, and so on. You need to flatten it out into a sequence of its scalar items (the leaves, if you think of the nested sequence as a tree).

Solution

Of course, we need to be able to tell which of the elements we’re handling are to be deemed scalar. For generality, say we’re passed as an argument a predicate that defines what is scalar—a function that we can call on any element and that returns 1 if the element is scalar or 0 otherwise. Given this, one approach is:

def flatten(sequence, scalarp, result=None):
    if result is None: result = []
    for item in sequence:
        if scalarp(item): result.append(item)
        else: flatten(item, scalarp, result)
    return result

In Python 2.2, a simple generator is an interesting alternative, and, if all the caller needs to do is loop over the flattened sequence, may save the memory needed for the result list:

from _ _future_ _ import generators
def flatten22(sequence, scalarp):
    for item in sequence:
        if scalarp(item):
            yield item
        else:
            for subitem in flatten22(item, scalarp):
                yield subitem

Discussion

The only problem with this recipe is that determining what is a scalar is not as obvious as it might seem, which is why I delegated that decision to a callable predicate argument that the caller is supposed to pass to flatten. Of course, we must be able to loop over the items of any non-scalar with a for statement, or flatten will raise an ...