Recursive list comprehension

Steven Bethard steven.bethard at gmail.com
Wed Dec 8 15:02:28 EST 2004


Adam DePrince wrote:
> def flatten( i ):
>     try:
>         i = i.__iter__()
>         while 1:
>             j = flatten( i.next() )
>             try:
>                 while 1:
>                     yield j.next()
>             except StopIteration:
>                 pass
>     except AttributeError:
>         yield i


Probably you want to catch a TypeError instead of an AttributeError; 
objects may support the iterator protocol without defining an __iter__ 
method:

 >>> class C:
...     def __getitem__(self, index):
...         if index > 3:
...             raise IndexError(index)
...         return index
...
 >>> list(iter(C()))
[0, 1, 2, 3]

I would write your code as something like:

 >>> def flatten(i):
...     try:
...         if isinstance(i, basestring):
...             raise TypeError('strings are atomic')
...         iterable = iter(i)
...     except TypeError:
...         yield i
...     else:
...         for item in iterable:
...             for sub_item in flatten(item):
...                 yield sub_item
...
 >>> list(flatten([['N','F'],['E'],['D']]))
['N', 'F', 'E', 'D']
 >>> list(flatten([C(), 'A', 'B', ['C', C()]]))
[0, 1, 2, 3, 'A', 'B', 'C', 0, 1, 2, 3]

Note that I special-case strings because, while strings support the 
iterator protocol, in this case we want to consider them 'atomic'.  By 
catching the TypeError instead of an AttributeError, I can support 
old-style iterators as well.

Steve



More information about the Python-list mailing list