itertools.flatten()? and copying generators/iterators.

[Peter Otten]

Raymond Hettinger wrote:

> Core itertools should be primitive building blocks that combine with one
> another to make other tools.  Also, I'm trying to keep the toolset as
> small as possible by excluding new tools that can be easily and
> efficiently coded in pure python.
> 
> I would code it like this:
> 
> def flatten(s):
>     try:
>         iter(s)
>     except TypeError:
>         yield s
>     else:
>         for elem in s:
>             for subelem in flatten(elem):
>                 yield subelem
> 
> As your examples show, it does have some suprising behavior in that
> strings get split apart instead of staying intact.  That is easily taken
> care of by an AtomicString subclass:
> 
> class AtomicString(str):
>     def __iter__(self):
>         raise TypeError
> 
> a = [1, 2, AtomicString('abc'), 4]
> 
> 
> 
>> >>> # The following I'm not sure what to do about...
>> >>> empty = [1, [], 3]
>> >>> emptyiter = [1, iter([]), 3]
> 

I suggest a minor modification:

def flatten(s, toiter=iter):
    try:
        it = toiter(s)
    except TypeError:
        yield s
    else:
        for elem in it:
            for subelem in flatten(elem, toiter):
                yield subelem

def keepstrings(seq):
    if isinstance(seq, basestring):
        raise TypeError
    return iter(seq)

sample = [1, 2, [3, "abc def".split()], 4]
print sample
print list(flatten(sample, keepstrings))
print list(flatten([1, [], 3]))
print list(flatten([1, iter([]), 3]))

The above handles strings in a way that is nonintrusive on client code.

Peter