4
\$\begingroup\$

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).

\$\endgroup\$
2
\$\begingroup\$

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
\$\endgroup\$
  • 4
    \$\begingroup\$ In Python 3.3+ you can replace the inner for loop on the recursive call with yield from flatten(item). \$\endgroup\$ – Blckknght May 15 '14 at 19:57
  • \$\begingroup\$ that was lazy enough ;-) \$\endgroup\$ – LetMeSOThat4U May 16 '14 at 13:06
1
\$\begingroup\$

You can use the compiler module

from compiler.ast import flatten

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

I hope this helps

\$\endgroup\$
  • 1
    \$\begingroup\$ 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) \$\endgroup\$ – LetMeSOThat4U May 16 '14 at 13:05

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.