Lazily transpose a list in Python

Question

So, I have an iterable of 3-tuples, generated lazily. I'm trying to figure out how to turn this into 3 iterables, consisting of the first, second, and third elements of the tuples, respectively. However, I wish this to be done lazily.

So, for example, I wish [(1, 2, 3), (4, 5, 6), (7, 8, 9)] to be turned into [1, 4, 7], [2, 5, 8], [3, 6, 9]. (Except I want iterables not lists.)

The standard zip(*data) idiom doesn't work, because the argument unpacking expands the entire iterable. (You can verify this by noting that zip(*((x, x+1, x+2) for x in itertools.count(step=3))) hangs.)

The best I've come up with thus far is the following:

def transpose(iterable_of_three_tuples):
    teed = itertools.tee(iterable_of_three_tuples, 3)
    return map(lambda e: e[0], teed[0]), map(lambda e: e[1], teed[1]), map(lambda e: e[2], teed[2])

This seems to work. But it hardly seems like clean code. And it does a lot of what seems to be unnecessary work.

Answer 1

Your transpose is pretty much exactly what you need.

With any solution you'd choose, you'd have to buffer the unused values (e.g. to get to the 7, you have to read 1-6, and store them in memory for when the other iterables ask for them). tee already does exactly that kind of buffering, so there's no need implementing it yourself.

The only other (minor) thing is that I'd write it slightly differently, avoiding the map and lambdas:

def transpose(iterable_of_three_tuples):
    teed = itertools.tee(iterable_of_three_tuples, 3)
    return ( e[0] for e in teed[0] ),  ( e[1] for e in teed[1] ),  ( e[2] for e in teed[2] )