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

Question

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

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

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

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

为了确保调用上下文管理器，我的 __next__ returns 包装器 class 不能直接替换值，因此如果用户忘记调用上下文管理器。

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

for x in some_iterator_object:
  ...

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

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

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

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

Answer 1

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

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 class 负责在返回值之前调用 __enter__ 其值。 __exit__ 在返回另一个值之前（如果一切顺利）或在循环中引发异常时调用。

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

Friendly usage of a Python iterable over a sequence of context managers

python

iterator

contextmanager

python-3.x