如何在Python中合并两个生成器(或其他可迭代对象)?

272

我想修改以下代码

for directory, dirs, files in os.walk(directory_1):
    do_something()

for directory, dirs, files in os.walk(directory_2):
    do_something()

对于这段代码:

for directory, dirs, files in os.walk(directory_1) + os.walk(directory_2):
    do_something()

我遇到了以下错误:

unsupported operand type(s) for +: 'generator' and 'generator'

如何在Python中合并两个生成器?

15个回答

358

itertools.chain()可以实现这个功能。它接受多个可迭代对象,并逐一从每个对象中产生元素,大致相当于:

def chain(*iterables):
    for it in iterables:
        for element in it:
            yield element

使用示例:

from itertools import chain

g = (c for c in 'ABC')  # Dummy generator, just for example
c = chain(g, 'DEF')  # Chain the generator and a string
for item in c:
    print(item)

输出:

A
B
C
D
E
F

15
需要记住的是,itertools.chain() 的返回值不会返回 types.GeneratorType 实例。如果确切的类型很关键,就要注意这一点。 - Riga
1
参见 @andrew-pate 的回答,了解 itertools.chain.from_iterable() 返回 types.GeneratorType 实例的相关信息。 - gkedge
itertools.chain()会给出一个目录中的所有元素,然后转移到另一个目录。 现在,我们如何选择两个目录的第一个元素并执行一些操作,然后转移到下一对,以此类推?任何想法都将不胜感激。 - yash
1
@yash 使用内置函数 next 手动迭代这些目录。 - Jeyekomon
2
@yash 你可能会喜欢 zip。它可以精确地选择第一个、第二个等值并将它们放入元组中。 - Randelung

114

一段代码示例:

from itertools import chain

def generator1():
    for item in 'abcdef':
        yield item

def generator2():
    for item in '123456':
        yield item

generator3 = chain(generator1(), generator2())
for item in generator3:
    print item

24
为什么不将这个示例添加到已经存在且高票的 itertools.chain() 答案中呢? - Jean-François Corbett
6
因为这会让他失去850点声望值,而他目前只有851分。你做你的就好了,cesio。 - Tatarize
2
@Jean-FrançoisCorbett写这个“已经存在”的答案的人本来就可以做到这一点...好吗? :) - Ice Bear
一个示例已经添加到顶部答案,这使得此内容变得多余。 - wjandrea

91

在Python(3.5或更高版本)中,您可以执行以下操作:

def concat(a, b):
    yield from a
    yield from b

16
太像Python风格了。 - ramazan polat
25
更通用的形式:def chain(*可迭代对象): for 可迭代对象中的每个元素 in 可迭代对象: yield from 元素 (在运行时将deffor放在不同的行上)。 - wjandrea
2
a 中的所有内容是否在 b 中的任何内容被生成之前生成,还是它们是交替生成的? - problemofficer - n.f. Monica
4
@problemofficer 是的。只有在从 a 中产生所有内容之后,才会检查它是否是迭代器,即使 b 不是迭代器也是如此。b 不是迭代器的 TypeError 将稍后出现。 - GeeTransit
1
@Karolius 噢,好的,我明白你的意思了。看起来你打错了一个字母,这让我很困惑: def chain(iterable) 应该是 def chain(iterables)。(另外,x for x in 是多余的。)无论如何,在标准库中已经有一个工具可以做到这一点:itertools.chain.from_iterable。除了性能之外,如果你有一个无限可迭代的可迭代对象,使用解包是不可能的 - wjandrea
显示剩余4条评论

41

简单示例:

from itertools import chain
x = iter([1,2,3])      #Create Generator Object (listiterator)
y = iter([3,4,5])      #another one
result = chain(x, y)   #Chained x and y

8
为什么不将这个例子加到已经存在且得到高票的 itertools.chain() 回答中呢? - Jean-François Corbett
这不太对,因为itertools.chain返回的是一个迭代器,而不是生成器。 - David J.
你不能只是这样做chain([1, 2, 3], [3, 4, 5])吗? - Corman
严谨地说,list_iterator 不是一个生成器,但它是一个迭代器,这实际上是 OP 所问的,因为在这种情况下,生成器与迭代器的行为没有任何不同。 - wjandrea
一个示例已经添加到顶部答案,这使得此内容变得多余。 - wjandrea

14
这里使用了一个带有嵌套的生成器表达式for循环。
range_a = range(3)
range_b = range(5)
result = ( item
           for one_range in (range_a, range_b)
           for item in one_range )
assert list(result) == [0, 1, 2, 0, 1, 2, 3, 4]
for ... in ... 是从左到右进行评估的。在 for 后面的标识符建立了一个新的变量。虽然在下面的 for ... in ... 中使用了 one_range,但是第二个中的 item 被用于“最终”赋值表达式,这个表达式只有一个(在一开始)。
相关问题:如何将列表的列表转换为平面列表?

13

使用itertools.chain.from_iterable,您可以执行以下操作:

def genny(start):
  for x in range(start, start+3):
    yield x

y = [1, 2]
ab = [o for o in itertools.chain.from_iterable(genny(x) for x in y)]
print(ab)

1
你正在使用不必要的列表推导式。当genny已经返回一个生成器时,你也在使用不必要的生成器表达式。list(itertools.chain.from_iterable(genny(x)))更加简洁。 - Corman
列表推导式是根据问题要求创建两个生成器的简便方式。也许我的回答在这方面有点复杂。 - andrew pate
1
我猜我添加这个答案到现有的答案中的原因是为了帮助那些需要处理大量生成器的人。 - andrew pate
Corman,我同意你的列表构造器确实更易读。不过,看到你的“许多更简单的方法”会很好……我认为wjandrea上面的评论似乎与itertools.chain.from_iterable做的一样,最好比赛一下,看看谁更快。 - andrew pate
如前所述,两种更简单的方法是使用listgenny(x)而不是列表推导式和生成器。速度竞赛几乎肯定会偏向于列表推导式,因为您要进行的计算较少。 - Corman
显示剩余2条评论

8

2020最新更新:适用于Python 3和Python 2

import itertools

iterA = range(10,15)
iterB = range(15,20)
iterC = range(20,25)

第一个选项

for i in itertools.chain(iterA, iterB, iterC):
    print(i)

# 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

这是Python 2.6中引入的备选选项

for i in itertools.chain.from_iterable( [iterA, iterB, iterC] ):
    print(i)

# 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

itertools.chain() 是基本方法。

如果你有可迭代的嵌套序列,例如一个子目录下的文件列表 [["src/server.py", "src/readme.txt"], ["test/test.py"]] ,那么itertools.chain.from_iterable() 就很有用了。


Python 2于2020年1月1日停止维护,因此我很惊讶你还提到它。 - wjandrea

3

我们也可以使用展开运算符*

concat = (*gen1(), *gen2())

注意:此方法最有效地适用于“非惰性”可迭代对象。也可与不同类型的推导式配合使用。对于生成器连接,最佳方式是使用@Uduse的答案。


很遗憾,*生成器没有惰性求值,因为这将使它成为一个奇妙的解决方案... - Camion
7
这将立即将两个生成器合并成一个元组! - wim

2

如果您想将生成器分开,但仍然同时迭代它们,可以使用zip():

注意:迭代会在两个生成器中较短的一个结束。

例如:

for (root1, dir1, files1), (root2, dir2, files2) in zip(os.walk(path1), os.walk(path2)):

    for file in files1:
        #do something with first list of files

    for file in files2:
        #do something with second list of files

2

免责声明:仅适用于Python 3!

语法类似于您想要的内容,可以使用星号运算符来展开这两个生成器:

for directory, dirs, files in (*os.walk(directory_1), *os.walk(directory_2)):
    do_something()

解释:

这实际上对两个生成器进行单层展平,得到一个由3元组构成的N元组(从os.walk中获取),形式如下:

((directory1, dirs1, files1), (directory2, dirs2, files2), ...)

你的for循环随后遍历这个N元组。

当然,只需用方括号替换外部括号,即可获得3元组列表,而不是3元组的N元组:

for directory, dirs, files in [*os.walk(directory_1), *os.walk(directory_2)]:
    do_something()

这将产生类似以下的内容:
[(directory1, dirs1, files1), (directory2, dirs2, files2), ...]

优点:

这种方法的好处是不需要导入任何东西,代码量也不多。

缺点:

缺点是你将两个生成器放入一个集合中,然后迭代该集合,实际上进行了两次遍历,可能会使用大量内存。


这根本不是扁平化。相反,它是一个 zip - jpaugh
2
有点困惑你的评论 @jpaugh。这个函数是将两个可迭代对象连接起来,而不是从它们中创建成对的元素。也许混淆是因为 os.walk 已经生成了 3 元组? - Milosz

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接