26

I have two generators say A() and B(). I want to iterate over both the generators together. Something like:

for a,b in A(),B():    # I know this is wrong
    #do processing on a and b

One way is to store the results of both the functions in lists and then loop over the merged list. Something like this:

resA = [a for a in A()]
resB = [b for b in B()]
for a,b in zip(resA, resB):
    #do stuff

If you are wondering, then yes both the functions yield equal number of value.

But I can't use this approach because A()/B() returns so many values. Storing them in a list would exhaust the memory, that's why I am using generators.

Is there any way to loop over both the generators at once?

Abhishek Gupta
  • 6,465
  • 10
  • 50
  • 82

3 Answers3

43

You were almost there. In Python 3, just pass the generators to zip():

for a, b in zip(A(), B()):

zip() takes any iterable, not just lists. It will consume the generators one by one.

In Python 2, use itertools.izip():

from itertools import izip

for a, b in izip(A(), B()):

As an aside, turning a generator into a list is as simple as list(generator); no need to use a list comprehension there.

Martijn Pieters
  • 1,048,767
  • 296
  • 4,058
  • 3,343
  • Purely academically, if `A()` yields 1,2,3 and `B()` yields 4,5,6 does this produce `a,b = (1,4)`, `(1,5)`, `(1,6)`, `(2,4)`, `(2,5)` ... `(3,6)` ? – Adam Smith Jan 03 '14 at 18:58
  • 3
    @adsmith: It does not. It produces `(1, 4)`, `(2, 5)` and `(3, 6)`. You are thinking of `itertools.product()`. – Martijn Pieters Jan 03 '14 at 18:59
3

Sounds like you want itertools.izip:

from itertools import izip

for a, b in izip(A(), B()):

From the docs:

Like zip() except that it returns an iterator instead of a list.

So this way you never create a list, either of A(), B() or the izip().

Note that in Python 3's basic zip is like Python 2.x's itertools.izip.

jonrsharpe
  • 115,751
  • 26
  • 228
  • 437
  • in this case both the lists has to have equal number of items, and if second list has more items it will not be printed – Ciasto piekarz May 20 '14 at 09:38
  • @san to be specific, whichever argument is longer will be truncated to the length of the shortest. If you need to process pairs, that is necessary. If you want to get as many items as possible, there is `[i]zip_longest`, although if one of your arguments is e.g. an infinitely long generator you will have a problem. – jonrsharpe May 20 '14 at 09:41
0

You can use the zip function to transform the two lists generated into an list of tuples, so you can use the way you want:

for a,b in zip(A(), B()):
    pass
Daniel César Indiano
  • 9
  • 3