超出范围时在生成器中进行清理

Question

我有一个生成器，即使它从未迭代过，它也必须执行清理步骤：

def gen(data):
    while True:
        item = data.get()
        if item is None:
            break
        # ...
        try:
            yield transformed_item
        except GeneratorExit:
            break
    # clean up; must happen if gen was called
    # ...

当我这样称呼它时，一切正常（即清理发生）：

for x in gen(data):
    # ...

或者像这样：

g = gen(data)
r = next(g)
# ...

但是当生成器超出范围而没有任何人对其调用 next 时，当然它根本不会执行任何代码，因此 GeneratorExit 不会在其中引发，并且清理不会发生：

g = gen(data)
# g was never used before going out of scope
del g

即使生成器在有机会产生任何东西之前就超出范围，我如何重构代码以确保执行清理步骤？

Answer 1

您可以为此使用上下文处理程序。这取决于您需要将生成器持久化多长时间。

class Gen(object):

    def __init__(self, data):
        self.data = data

    def __enter__(self):
        return self._gen(self.data)

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Cleanup
        print 'Cleaning up'

    def _gen(self, data):
        for i in data:
            yield i

然后它看起来像：

with Gen(data) as g:
    r = next(g)

编辑：

鉴于您不能强制最终用户使用上下文管理器的限制，您能否将生成器创建包装在另一个函数中并“播种”生成器？

def gen(data):
    g = _gen(data)
    next(g)
    return g


def _gen(data):
    yield None
    while True:
        ... # Rest of generator