Nisba
Nisba

Reputation: 3448

Python nested generators not working

I am practising with generators and I wonder why the following code does not print 16 pairs but only 4.

def range_generator_function(my_range):
    for i in my_range:
        yield i

gen1=range_generator_function(range(1,5))
gen2=range_generator_function(range(1,5))

def pairs_generator_function(gen1,gen2):
    for it1 in gen1:
        for it2 in gen2:
            yield [it1,it2]

my_gen = pairs_generator_function(gen1,gen2)

for it in my_gen:
    print(it)

The output is

[1, 1]
[1, 2]
[1, 3]
[1, 4]

While the output I expect is

[1, 1]
[1, 2]
[1, 3]
[1, 4]
[2, 1]
[2, 2]
[2, 3]
[2, 4]
[3, 1]
[3, 2]
[3, 3]
[3, 4]
[4, 1]
[4, 2]
[4, 3]
[4, 4]

Upvotes: 3

Views: 343

Answers (5)

Olivier Melançon
Olivier Melançon

Reputation: 22324

The output you get is due to the fact that on the second iteration, the generator gen2 has been exhausted.

You can either make a list out of it to store its output prior to looping or copy it with itertools.tee on every iteration. The former will not be able to handle infinite generators, but before implementing the later let me point out that you are actually reimplementing itertools.product.

from itertools import product

my_gen = product(gen1, gen2)

for it in my_gen:
    print(it)

Upvotes: 1

Ajax1234
Ajax1234

Reputation: 71471

As @wim pointed out, your generator is entirely exhausted after the first instance of iteration. However, to prevent this, cast gen1 and gen2 as lists when passing them to pairs_generator_function. However, you can use itertools.tee to store two copies of the original generator: one to pass to the function as a list, and the other for future use:

import itertools

def range_generator_function(my_range):
   for i in my_range:
     yield i

gen1=range_generator_function(range(1,5))
gen2=range_generator_function(range(1,5))

def pairs_generator_function(gen1,gen2):
  for it1 in gen1:
    for it2 in gen2:
        yield [it1,it2]

gen1, gen1_l = itertools.tee(gen1)
gen2, gen2_l = itertools.tee(gen2)
print(list(pairs_generator_function(list(gen1_l), list(gen2_l))))

Output:

[[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [2, 3], [2, 4], [3, 1], [3, 2], [3, 3], [3, 4], [4, 1], [4, 2], [4, 3], [4, 4]]

Notice, however, that gen1 and gen2` still point to the items in memory:

>>next(gen1)
1
>>next(gen2)
1

Upvotes: 0

Nisba
Nisba

Reputation: 3448

I modified the code so I get the wanted output, now the generator is recreated every time inside the outer loop.

def range_generator_function(my_range):
    for i in my_range:
        yield i

def pairs_generator_function():
    gen1=range_generator_function(range(1,5))
    for it1 in gen1:
        gen2=range_generator_function(range(1,5))
        for it2 in gen2:
            yield [it1,it2]

my_gen = pairs_generator_function()

for it in my_gen:
    print(it)

Upvotes: 0

Prune
Prune

Reputation: 77900

You exhausted gen2 after the first time through the inner loop. You no longer have anything in that stream, so the other three values of it1 have nothing with which to pair. You need to restart gen2 every time through. Consider itertools.tee for cloning more copies.

Upvotes: 2

wim
wim

Reputation: 363476

The actual output is correct. Your gen2 instance is completely exhausted by the first inner loop:

def pairs_generator_function(gen1,gen2):
    for it1 in gen1:
        for it2 in gen2:  # <--- this consumes gen2
            yield [it1,it2]

On the subsequent iterations, iterating gen2 again is empty.

Upvotes: 5

Related Questions