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

• Blckknght: In Python 3.3+ you can replace the inner for loop on the recursive call with yield from flatten(item).
• LetMeSOThat4U: that was lazy enough ;-)

Answer 2

You can use the compiler module

from compiler.ast import flatten
>>> flatten([1,[2,3])
>>> [1,2,3]

I hope this helps

• LetMeSOThat4U: look at /lib/python2.7/compiler/ast.py / flatten definition, that's plain eager evaluation flatten like one I was trying to rewrite, while I'm looking for lazy iterator (lazy ~= conserves memory)