Smoothing Nested Generator Expressions

Question

I am trying to smooth out a nested generator generator, but I get an unexpected result:

>>> g = ((3*i + j for j in range(3)) for i in range(3)) >>> list(itertools.chain(*g)) [6, 7, 8, 6, 7, 8, 6, 7, 8] 

I expected the result to look like this:

 [0, 1, 2, 3, 4, 5, 6, 7, 8] 

I think I get an unexpected result, because the internal generators are not evaluated until the external generator is repeated, setting i to 2. I can crack the solution by forcing the evaluation of internal generators using a list expression instead of a generator expression:

 >>> g = ([3*i + j for j in range(3)] for i in range(3)) >>> list(itertools.chain(*g)) [0, 1, 2, 3, 4, 5, 6, 7, 8] 

Ideally, I would like the solution to be completely lazy and not forcibly evaluate internal nested elements until they are used.

Is there a way to smooth out nested expressions of a generator of arbitrary depth (possibly using something other than itertools.chain )?

Edit:

No, my question is not a duplicate of Variable Scope In Generators In Classes . I honestly cannot say how these two issues are related at all. Maybe the moderator could explain why he thinks this is a duplicate.

In addition, both answers to my question are correct in that they can be used to write a function that correctly aligns nested generators.

 def flattened1(iterable): iter1, iter2 = itertools.tee(iterable) if isinstance(next(iter1), collections.Iterable): return flattened1(x for y in iter2 for x in y) else: return iter2 def flattened2(iterable): iter1, iter2 = itertools.tee(iterable) if isinstance(next(iter1), collections.Iterable): return flattened2(itertools.chain.from_iterable(iter2)) else: return iter2 

As far as I can tell with timeit , they both execute the same way.

 >>> timeit(test1, setup1, number=1000000) 18.173431718023494 >>> timeit(test2, setup2, number=1000000) 17.854709611972794 

I'm not sure which one is better in terms of style, since x for y in iter2 for x in y is a bit of a brain twister, but maybe more elegant than itertools.chain.from_iterable(iter2) . Entrance is appreciated.

Unfortunately, I could only point out one of two equally good answers.

Answer 1

Instead of chain(*g) you can use chain.from_iterable :

 >>> g = ((3*i + j for j in range(3)) for i in range(3)) >>> list(itertools.chain(*g)) [6, 7, 8, 6, 7, 8, 6, 7, 8] >>> g = ((3*i + j for j in range(3)) for i in range(3)) >>> list(itertools.chain.from_iterable(g)) [0, 1, 2, 3, 4, 5, 6, 7, 8] 

Answer 2

How about this:

 [x for y in g for x in y] 

What gives:

 [0, 1, 2, 3, 4, 5, 6, 7, 8] 

Answer 3

Guess that you already have your answer, but here is a different perspective.

The problem is that when each internal generator is created, the expression generating the value is closed by the external variable i , so even when the first internal generator starts to generate values, it uses the "current" value of i . This value will have the value i=2 if the external generator is completely consumed (and this is exactly the same as after the argument in the chain(*g) call before chain actually called).

The following problem will work around the problem:

 g = ((3*i1 + j for i1 in [i] for j in range(3)) for i in range(3)) 

Note that these internal generators are not closed over i because the for clauses are evaluated at the time the generator was created, so a singleton list [i] computed and its value is “frozen” in the face of further changes to the i value.

This approach has the advantage over the from_iterable answer that it is a little more general if you want to use it outside the chain.from_iterable call - it will always generate the “correct” internal generators, regardless of whether the external generator is partially or fully consumed before using the internal generators. For example, in the following code:

 g = ((3*i1 + j for i1 in [i] for j in range(3)) for i in range(3)) g1 = next(g) g2 = next(g) g3 = next(g) 

you can insert lines:

 list(g1) list(g2) list(g3) 

in any order at any point after determining the appropriate internal generator, and you will get the correct results.