Search code examples
pythonfor-loopwhile-loopiteratorpython-itertools

Nesting iteration - difference between for and while loops


I need to do a nested iteration over a generator (not a list). What I need is to perform something like this:

testing  3 ...
Testing passed!
     Starting subtest:
     Sub-testing 4  with  3
     Sub-testing passed!
testing  4 ...
testing  5 ...
testing  6 ...
Testing passed!
     Starting subtest:
     Sub-testing 7  with  6
     Sub-testing 8  with  6
     Sub-testing 9  with  6
     Sub-testing passed!
testing  7 ...
testing  8 ...
testing  9 ...
Testing passed!
     Starting subtest:
     Sub-testing 10  with  9
     Sub-testing 11  with  9
     Sub-testing 12  with  9
     Sub-testing passed!
testing  10 ...

So I tried the following code , using a for loop:

from itertools import *
princ_iter = count(3)
for x in princ_iter:
    print("testing ", x, "...")
    if x % 3 == 0:
        print("Testing passed!")
        print("     Starting subtest:")
        princ_iter, nested_iter = tee(princ_iter)
        for y in nested_iter:
            print("     Sub-testing", y, " with ", x)
            if y % (x//2) == 0:
                print("     Sub-testing passed!")
                break

but it does not work because the principal iterator (princ_iter) is iterated together with the nested one (nested_iter) and I obtain this output instead:

testing  3 ...
Testing passed!
     Starting subtest:
     Sub-testing 4  with  3
     Sub-testing passed!
testing  5 ...
testing  6 ...
Testing passed!
     Starting subtest:
     Sub-testing 4  with  6
     Sub-testing 7  with  6
     Sub-testing 8  with  6
     Sub-testing 9  with  6
     Sub-testing passed!
testing  10 ...
testing  11 ...

So I tried using the same instructions in a while loop:

from itertools import *
princ_iter= count(3)
while True:
    x = next(princ_iter)
    print("testing ", x, "...")
...

and this time I have obtained exactly the output I was looking for !

Why there is this difference between the two instructions? Is there a (better) way to do it using a for loop?


Solution

  • This is behavior is alluded to in the documentations for tee:

    Once tee() has made a split, the original iterable should not be used anywhere else; otherwise, the iterable could get advanced without the tee objects being informed.

    When you use a for-loop, the original iterator is being used the whole time:

    for x in princ_iter:
        ...
    

    The for-loop will always work with the same object.

    The fact that:

    princ_iter, nested_iter = tee(princ_iter)
    

    Re-assigns the prince_iter variable with a new iterator is irrelevant.

    On the other hand, in your while-loop, this is relevant, because you are controlling what iterator is advanced:

    x = next(princ_iter)
    

    i.e., whatever iterator is being currently referred to by prince_iter, so the variable re-assignment does affect things.