Recursive flatten with lazy evaluation iterator

Question

I'm trying to rewrite this:

def flatten(lst):
    flat = []
    for x in lst:
        if hasattr(x, '__iter__') and not isinstance(x, basestring):
            flat.extend(flatten(x))
        else:
            flat.append(x)
    return flat


In [21]: a=[1, [2, 3, 4, [5, 6]], 7]

In [22]: flatten(a)
Out[22]: [1, 2, 3, 4, 5, 6, 7]

..into a version that would flatten the iterable, but return values in a lazy manner. Now, this works:

def flat_fugly(s):
    if iterable(s):
        for x in s:
            yield chain.from_iterable(flat_fugly(x))
    else:
        yield takewhile(lambda x: True, [s])


list(islice(chain.from_iterable(flat_fugly(a)), 0, 6))
Out[34]: [1, 2, 3, 4, 5, 6]

But, as the name suggests... Is there a cleaner/better way?

Not to mention that I have to apply chain.from_iterable to flat_fugly anyway while I'd prefer to have plain iterator (I could wrap it in another function that would use chain.from_iterable of course, but still if it could be all made to fit in one more elegant function that would be preferable).

Answer 1

Not my idea, but I think this is better:

import collections


def flatten(lst):
    for item in lst:
        if isinstance(item, collections.Iterable) and not isinstance(item, basestring):
            for sublst in flatten(item):
                yield sublst
        else:
            yield item

Answer 2

You can use the compiler module

from compiler.ast import flatten

>>> flatten([1,[2,3])
>>> [1,2,3]

I hope this helps