python - 展平嵌套的生成器表达式

Question

我正在尝试展平一个嵌套的生成器生成器，但我得到了一个意想不到的结果:

>>> 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]

我希望结果看起来像这样:

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

我认为我得到了意想不到的结果，因为在外部生成器已经迭代之前不会评估内部生成器，将 i 设置为 2。我可以通过强制通过使用列表理解而不是生成器表达式来评估内部生成器:

>>> 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]

理想情况下，我想要一个完全惰性的解决方案，并且在使用内部嵌套元素之前不会强制计算它们。

有没有办法扁平化任意深度的嵌套生成器表达式(可能使用 itertools.chain 以外的东西)？

编辑:

不，我的问题不是 Variable Scope In Generators In Classes 的重复问题.老实说，我根本不知道这两个问题是如何相关的。也许主持人可以解释为什么他认为这是重复的。

另外，我的问题的两个答案都是正确的，因为它们可以用来编写一个函数来正确地展平嵌套生成器。

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

据我所知，timeit，它们的表现是一样的。

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

我也不确定从风格的角度来看哪个更好，因为 x for y in iter2 for x in y 有点绕口令，但可以说比 itertools.chain.from_iterable(iter2)。感谢输入。

遗憾的是，我只能将两个同样好的答案中的一个标记为正确。

Answer 1

您可以使用 chain.from_iterable 而不是使用 chain(*g) :

>>> 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]