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

Question

我正在编写一个小库，试图为调度作业提供一个持久的队列。我的持久性代码提供了一种迭代挂起作业描述的方法;我还想保证已调度的作业最终被标记为已完成或失败。

为此，我首先实现了它，以便我的用户可以执行以下操作：

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

出于多种原因，我不喜欢这种解决方案。首先，我想从我的队列中获取作业描述作为单个操作(如果队列为空，则失败(，因此获取是通过迭代器的__next__方法完成的，并在上下文管理器的__exit__中发布。

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

有没有办法将这两个语句折叠成一个？理想情况下，我想让用户做

for x in some_iterator_object:
...

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

编辑：我通过实验发现，如果我让一个未完成的生成器被垃圾收集，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__在返回另一个值之前(如果一切顺利(或在循环中引发异常时调用。