Split iterator into multiple streams

[Raymond Hettinger]

On Nov 6, 1:52 am, Steven D'Aprano <st... at REMOVE-THIS-
cybersource.com.au> wrote:
> I tried changing the t[i] to use operator.itergetter instead, but no
> luck. Finally I got this:
>
> def split(iterable, n):
>     iterators = []
>     for i, iterator in enumerate(itertools.tee(iterable, n)):
>         f = lambda it, i=i: (t[i] for t in it)
>         iterators.append(f(iterator))
>     return tuple(iterators)
>
> which seems to work:
>
> >>> data = [(1,2,3), (4,5,6), (7,8,9)]
> >>> a, b, c = split(data, 3)
> >>> list(a), list(b), list(c)
>
> ([1, 4, 7], [2, 5, 8], [3, 6, 9])
>
> Is this the right approach, or have I missed something obvious?


That looks about right to me.
It can be compacted a bit:

def split(iterable, n):
  return tuple(imap(itemgetter(i), it) for i, it in
enumerate(tee(iterable, n)))

Internally, the tee's iterators are going to accumulate a ton of data
unless they are consumed roughly in parallel.  Of course, if they are
consumed *exactly* in lockstep, the you don't need to split them into
separate iterables -- just use the tuples as they come.


Raymond