Iterable Flattener with Depth.

[Ian Clark]

thebjorn wrote:
> On Nov 2, 6:32 am, praddy <praddyjin... at gmail.com> wrote:
>> On Nov 1, 5:03 pm, bearophileH... at lycos.com wrote:
>>
>>> Pradeep Jindal:
>>>> Any comments?
>>> Something with similar functionality (plus another 20 utility
>>> functions/classes or so) has probably to go into the std lib... :-)
>>> Bye,
>>> bearophile
>> Same Here!
>>
>> - Pradeep
> 
> Yeah, everyone has to write a flatten sooner or later :-)  My version
> is at:
> 
>   http://blog.tkbe.org/archive/python-flatten/
> 
> -- bjorn
> 

And here is mine. Note that it is very similar to Michael Spencer's 
implementation[1]. The only difference is that this adds a depth counter.

def iflat(itr, depth=0):
   itr = iter(itr)
   stack = []
   cur_depth = 0

   while True:
     try:
       elem = itr.next()
       if hasattr(elem, "__iter__") and cur_depth < depth:
         stack.append(itr)
         itr = iter(elem)
         cur_depth += 1
       else:
         yield elem
     except StopIteration:
       if not stack:
         raise StopIteration
       cur_depth -= 1
       itr = stack.pop()


if __name__ == "__main__":
   test1 = ((0, 1, 2), ((3, 4), 5), (((6, 7), 8), 9))
   test2 = [1,[2,[3,4,5],'bash'],6,[7,[8,[9,10,['hi', 'hello']]]], 11, 12]

   for x in (test1, test2):
     print
     print list(iflat(x))
     print
     print list(iflat(x, 1))
     print list(iflat(x, 2))
     print list(iflat(x, 3))
     print list(iflat(x, 4))
     print iflat(x, 10)

Ian

[1] http://mail.python.org/pipermail/python-list/2005-March/312022.html