如何克隆 Python 生成器对象？

Question

考虑这种情况：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

walk = os.walk('/home')

for root, dirs, files in walk:
    for pathname in dirs+files:
        print os.path.join(root, pathname)

for root, dirs, files in walk:
    for pathname in dirs+files:
        print os.path.join(root, pathname)

我知道这个例子有点多余，但你应该考虑到我们需要多次使用相同的walk 数据。我有一个基准场景，并且必须使用相同的 walk 数据才能获得有用的结果。

我已经尝试walk2 = walk 在第二次迭代中克隆和使用，但没有成功。问题是......我怎样才能复制它？有可能吗？

提前谢谢你。

Answer 1

你可以使用itertools.tee():

walk, walk2 = itertools.tee(walk)

请注意，正如文档指出的那样，这可能“需要大量额外的存储空间”。

Answer 2

如果您知道每次使用都要遍历整个生成器，那么通过将生成器展开到一个列表并多次使用该列表，您可能会获得最佳性能。

walk = list(os.walk('/home'))

Answer 3

定义一个函数

 def walk_home():
     for r in os.walk('/home'):
         yield r

甚至这个

def walk_home():
    return os.walk('/home')

两者都是这样使用的：

for root, dirs, files in walk_home():
    for pathname in dirs+files:
        print os.path.join(root, pathname)

Answer 4

这是functools.partial() 的一个很好的用例 制作一个快速的发电机工厂：

from functools import partial
import os

walk_factory = partial(os.walk, '/home')

walk1, walk2, walk3 = walk_factory(), walk_factory(), walk_factory()

functools.partial() 所做的事情很难用人类语言来描述，但这^ 就是它的用途。

它部分填写函数参数而不执行该函数。因此，它充当函数/生成器工厂。

Answer 5

此答案旨在扩展/详细说明其他答案所表达的内容。解决方案必然会根据您的目标确切而有所不同。

如果您想多次迭代os.walk 的完全相同的结果，您需要从os.walk 可迭代的项目（即walk = list(os.walk(path))）中初始化一个列表。

如果您必须保证数据保持不变，那可能是您唯一的选择。但是，在某些情况下这是不可能或不可取的。

1. 如果输出足够大，则无法 list() 迭代（即尝试 list() 整个文件系统可能会冻结您的计算机）。
2. 如果您希望在每次使用前获取“新”数据，则不希望 list() 可迭代。

如果list() 不合适，您将需要按需运行您的生成器。请注意，发电机在每次使用后都会熄灭，所以这会带来一个小问题。为了多次“重新运行”您的生成器，您可以使用以下模式：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

class WalkMaker:
    def __init__(self, path):
        self.path = path
    def __iter__(self):
        for root, dirs, files in os.walk(self.path):
            for pathname in dirs + files:
                yield os.path.join(root, pathname)

walk = WalkMaker('/home')

for path in walk:
    pass

# do something...

for path in walk:
    pass

上述设计模式可以让您的代码保持干燥。

Answer 6

此“Python 生成器侦听器”代码允许您在单个生成器上拥有多个侦听器，例如 os.walk，甚至稍后有人“插话”。

def walkme(): os.walk('/home')

m1 = Muxer(walkme) m2 = Muxer(walkme)

那么 m1 和 m2 甚至可以在线程中运行并在空闲时处理。

见：https://gist.github.com/earonesty/cafa4626a2def6766acf5098331157b3

import queue
from threading import Lock
from collections import namedtuple

class Muxer():
    Entry = namedtuple('Entry', 'genref listeners, lock')

    already = {}
    top_lock = Lock()

    def __init__(self, func, restart=False):
        self.restart = restart
        self.func = func
        self.queue = queue.Queue()

        with self.top_lock:
            if func not in self.already:
                self.already[func] = self.Entry([func()], [], Lock())
            ent = self.already[func]

        self.genref = ent.genref
        self.lock = ent.lock
        self.listeners = ent.listeners

        self.listeners.append(self)

    def __iter__(self):
        return self

    def __next__(self):
        try:
            e = self.queue.get_nowait()
        except queue.Empty:
            with self.lock:
                try:
                    e = self.queue.get_nowait()
                except queue.Empty:
                    try:
                        e = next(self.genref[0])
                        for other in self.listeners:
                            if not other is self:
                                other.queue.put(e)
                    except StopIteration:
                        if self.restart:
                            self.genref[0] = self.func()
                        raise
        return e

    def __del__(self):
        with self.top_lock:
            try:
                self.listeners.remove(self)
            except ValueError:
                pass
            if not self.listeners and self.func in self.already:
                del self.already[self.func]