python - 在一系列上下文管理器上友好地使用 Python 可迭代对象

Question

我正在编写一个小型库，它试图为调度作业提供一个持久队列。我的持久性代码提供了一种迭代未决职位描述的方法；我还想保证分派(dispatch)的作业最终会被标记为已完成或失败。

为此，我首先实现了它，以便我的用户可以:

for c in some_iterator_object:
  with c as x:
    ...

我不喜欢这个解决方案有几个原因。首先，我想从我的队列中获取一个作业描述作为一个操作(如果队列为空则失败)，所以获取是通过迭代器的 __next__ 方法完成的，并且在上下文管理器的 __exit__ 中释放。

为了确保调用上下文管理器，我的 __next__ 返回一个不能直接替换值的包装类，因此如果用户忘记调用上下文管理器，它将抛出一个明确的错误.

有没有办法将这两个语句合并为一个语句？理想情况下，我想让用户做

for x in some_iterator_object:
  ...

同时能够拦截由 for block 的内容引发的异常。

编辑:我通过实验发现，如果我让一个未完成的生成器被垃圾收集，yield 语句将引发一个内部异常，所以我可以写一些粗糙的东西，比如

try:
  ...
  success = False
  yield val
  success = True
  ...
finally:
  if success:
     ...

但如果我理解正确的话，这取决于垃圾收集器来运行，而且它似乎是一种我不应该真正接触的内部机制。

Answer 1

如果您希望上下文管理器在迭代器返回时自动输入，您可以像这样编写自己的迭代器类:

class ContextManagersIterator:

    def __init__(self, it):
        self._it = iter(it)
        self._last = None

    def __iter__(self):
        return self

    def __next__(self):
        self.__exit__(None, None, None)

        item = next(self._it)
        item.__enter__()
        self._last = item

        return item

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        last = self._last
        if last is not None:
            self._last = None
            return last.__exit__(exc_type, exc_value, exc_traceback)

示例用法:

from contextlib import contextmanager

@contextmanager
def my_context_manager(name):
    print('enter', name)
    try:
        yield
    finally:
        print('exit', name)

sequence = [
    my_context_manager('x'),
    my_context_manager('y'),
    my_context_manager('z'),
]

with ContextManagersIterator(sequence) as it:
    for item in it:
        print('  work')

# Output:
# enter x
#   work
# exit x
# enter y
#   work
# exit y
# enter z
#   work
# exit z

这个 ContextManagersIterator 类负责在返回值之前对其值调用 __enter__。 __exit__ 在返回另一个值之前(如果一切顺利)或在循环中引发异常时调用。