Fast flatten() with depth control and oversight over which subtrees to expand (Python recipe) by Kevin L. Sitze

ActiveState Code (http://code.activestate.com/recipes/577470/)

#! /usr/bin/env python
######################################################################
#  Written by Kevin L. Sitze on 2010-11-25
#  This code may be used pursuant to the MIT License.
######################################################################

"""This module contains four flatten() functions for generating either
a sequence or an iterator.  The goal is to "flatten" a tree (typically
comprised of lists and tuples) by slicing the elements in each
contained subtree over the subtree element itself.  For example:

    ([a, b], [c, (d, e)]) => (a, b, c, d, e)

The functions available via this module are

    flatten    ( sequence[, max_depth[, ltypes]] ) => sequence
    xflatten   ( sequence[, max_depth[, ltypes]] ) => iterator
    flatten_it ( iterable[, max_depth[, ltypes]] ) => sequence
    xflatten_it( iterable[, max_depth[, ltypes]] ) => iterator

Each function takes as its only required argument the tree to flatten.
The first two functions (flatten() and xflatten()) require their first
argument to be a valid Python sequence object.  The '_it' functions
(flatten_it() and xflatten_it()) accept any iterable object.

The return type for the flatten() and xflatten_it functions is the
same type as the input sequence, when possible, otherwise the type
will be 'list'.

Wall clock speed of these functions increase from the top of the list
down (i.e., where possible prefer the flatten() function to any other
if speed is a concern).

The "max_depth" argument is either a non-negative integer indicating
the maximum tree depth to flatten; or "None" to flatten the entire
tree.  The required "sequence" argument has a 'depth' of zero; a list
element of the sequence would be flattened if "max_depth" is greater
than or equal to one (1).  A negative depth is treated the same as a
depth of zero.

The "ltypes" argument indicates which elements are subtrees.  It may
be either a collection of sequence types or a predicate function.  If
"ltypes" is a collection a subtree is expanded if its type is in the
collection.  If "ltypes" is a predicate function a subtree is expanded
if the predicate function returns "True" for that subtree.

The implementation of flatten here runs in O(N) time where N is the
number of elements in the traversed tree.  It uses O(N+D) space
where D is the maximum depth of the traversed tree.
"""

def flatten( l, max_depth = None, ltypes = ( list, tuple ) ):
    """flatten( sequence[, max_depth[, ltypes]] ) => sequence

    Flatten every sequence in "l" whose type is contained in "ltypes"
    to "max_depth" levels down the tree.  See the module documentation
    for a complete description of this function.

    The sequence returned has the same type as the input sequence.
    """
    if max_depth is None: make_flat = lambda x: True
    else: make_flat = lambda x: max_depth > len( x )
    if callable( ltypes ): is_sequence = ltypes
    else: is_sequence = lambda x: isinstance( x, ltypes )

    r = list()
    s = list()
    s.append(( 0, l ))
    while s:
        i, l = s.pop()
        while i < len( l ):
            while is_sequence( l[i] ):
                if not l[i]: break
                elif make_flat( s ):
                    s.append(( i + 1, l ))
                    l = l[i]
                    i = 0
                else:
                    r.append( l[i] )
                    break
            else: r.append( l[i] )
            i += 1
    try: return type(l)(r)
    except TypeError: return r

def xflatten( l, max_depth = None, ltypes = ( list, tuple ) ):
    """xflatten( sequence[, max_depth[, ltypes]] ) => iterable

    Flatten every sequence in "l" whose type is contained in "ltypes"
    to "max_depth" levels down the tree.  See the module documentation
    for a complete description of this function.

    This is the iterator version of the flatten function.
    """
    if max_depth is None: make_flat = lambda x: True
    else: make_flat = lambda x: max_depth > len( x )
    if callable( ltypes ): is_sequence = ltypes
    else: is_sequence = lambda x: isinstance( x, ltypes )

    r = list()
    s = list()
    s.append(( 0, l ))
    while s:
        i, l = s.pop()
        while i < len( l ):
            while is_sequence( l[i] ):
                if not l[i]: break
                elif make_flat( s ):
                    s.append(( i + 1, l ))
                    l = l[i]
                    i = 0
                else:
                    yield l[i]
                    break
            else: yield l[i]
            i += 1

def flatten_it( l, max_depth = None, ltypes = ( list, tuple ) ):
    """flatten_it( iterator[, max_depth[, ltypes]] ) => sequence

    Flatten every sequence in "l" whose type is contained in "ltypes"
    to "max_depth" levels down the tree.  See the module documentation
    for a complete description of this function.

    The sequence returned has the same type as the input sequence.
    """
    if max_depth is None: make_flat = lambda x: True
    else: make_flat = lambda x: max_depth > len( x )
    if callable( ltypes ): is_iterable = ltypes
    else: is_iterable = lambda x: isinstance( x, ltypes )

    r = list()
    s = list()
    s.append(( iter( l ) ))
    while s:
        i = s.pop()
        try:
            while True:
                e = i.next()
                if is_iterable( e ):
                    if make_flat( s ):
                        s.append(( i ))
                        i = iter( e )
                    else:
                        r.append( e )
                else:
                    r.append( e )
        except StopIteration: pass
    try: return type(l)(r)
    except TypeError: return r

def xflatten_it( l, max_depth = None, ltypes = ( list, tuple ) ):
    """xflatten_it( iterator[, max_depth[, ltypes]] ) => iterator

    Flatten every sequence in "l" whose type is contained in "ltypes"
    to "max_depth" levels down the tree.  See the module documentation
    for a complete description of this function.

    This is the iterator version of the flatten_it function.
    """
    if max_depth is None: make_flat = lambda x: True
    else: make_flat = lambda x: max_depth > len( x )
    if callable( ltypes ): is_iterable = ltypes
    else: is_iterable = lambda x: isinstance( x, ltypes )

    r = list()
    s = list()
    s.append(( iter( l ) ))
    while s:
        i = s.pop()
        try:
            while True:
                e = i.next()
                if is_iterable( e ):
                    if make_flat( s ):
                        s.append(( i ))
                        i = iter( e )
                    else:
                        yield e
                else:
                    yield e
        except StopIteration: pass

if __name__ == '__main__':

    import sys
    import traceback
    def assertEquals( exp, got ):
        if exp is got:
            r = True
        elif type( exp ) is not type( got ):
            r = False
        elif type( exp ) in ( float, complex ):
            r = abs( exp - got ) < 1e-8
        else:
            r = ( exp == got )
        if not r:
            print >>sys.stderr, "Error: expected <%s> but got <%s>" % ( repr( exp ), repr( got ) )
            traceback.print_stack()

    def test( exp, got, depth = None ):
        assertEquals( exp, flatten( got, depth ) )
        assertEquals( exp, tuple( xflatten( got, depth ) ) )
        assertEquals( exp, flatten_it( got, depth ) )
        assertEquals( exp, tuple( xflatten_it( got, depth ) ) )

    test( (),      () )
    test( (),      (()) )
    test( (),      ((),()) )
    test( (),      ((),((),()),()) )
    test( (1,),    ((1,),((),()),()) )
    test( (1,),    ((),1,((),()),()) )
    test( (1,),    ((),(1,(),()),()) )
    test( (1,),    ((),((1,),()),()) )
    test( (1,),    ((),((),1,()),()) )
    test( (1,),    ((),((),(1,)),()) )
    test( (1,),    ((),((),(),1),()) )
    test( (1,),    ((),((),()),1,()) )
    test( (1,),    ((),((),()),(1,)) )
    test( (1,),    ((),((),()),(),1) )
    test( (1,),    ((),1,()) )
    test( (1,2,3), (1,2,3) )
    test( (1,2,3), ((1,2),3) )
    test( (1,2,3), (1,(2,3)) )
    test( (1,2,3), ((1,),(2,),3) )
    test( ((((((((((0,),1),2),3),4),5),6),7),8),9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 0 )
    test( (((((((((0,),1),2),3),4),5),6),7),8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 1 )
    test( ((((((((0,),1),2),3),4),5),6),7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 2 )
    test( (((((((0,),1),2),3),4),5),6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 3 )
    test( ((((((0,),1),2),3),4),5,6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 4 )
    test( (((((0,),1),2),3),4,5,6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 5 )
    test( ((((0,),1),2),3,4,5,6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 6 )
    test( (((0,),1),2,3,4,5,6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 7 )
    test( ((0,),1,2,3,4,5,6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 8 )
    test( (0,1,2,3,4,5,6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 9 )
    test( (0,1,2,3,4,5,6,7,8,9), ((((((((((0,),1),2),3),4),5),6),7),8),9), 10 )

    test( ({1:2},3,4,set([5,6])), ({1:2},(3,4),set([5,6])) )

    l = (1,)
    # Build a tree 1 million elements deep
    for i in xrange( 1000000 ): l = ( l, 2 )
    # expected value is a 1 followed by 1 million 2's
    exp = (1,) + (2,) * 1000000
    # # Under 5 seconds on my machine...
    # got = flatten( l )
    # assert( exp == got )
    # # Also under 5 seconds...
    # got = tuple( xflatten( l ) )
    # assert( exp == got )
    # # 6 seconds
    # got = flatten_it( l )
    # assert( exp == got )
    # # 7 seconds
    # got = tuple( xflatten_it( l ) )
    # assert( exp == got )