如何制作一个内部有循环的上下文管理器？

Question

我想要这样的东西：

from contextlib import contextmanager

@contextmanager
def loop(seq):
    for i in seq:
        try:
            do_setup(i)
            yield # with body executes here
            do_cleanup(i)
        except CustomError as e:
            print(e)

with loop([1,2,3]):
    do_something_else()
    do_whatever()

但是 contextmanager 不起作用，因为它希望生成器只产生一次。

我想要这个的原因是因为我基本上想制作自己的自定义 for 循环。我有一个修改过的 IPython 用来控制测试设备。它显然是一个完整的 Python REPL，但大多数时候用户只是调用预定义的函数（类似于 Bash 提示符），并且用户不应是程序员或熟悉 Python.需要有一种方法可以使用 setup/cleanup 和每次迭代的异常处理来循环一些任意代码，并且它应该像上面的 with 语句一样简单。

Answer 1

我认为生成器在这里工作得更好：

def loop(seq):
    for i in seq:
        try:
            print('before')
            yield i  # with body executes here
            print('after')
        except CustomError as e:
            print(e)

for i in loop([1,2,3]):
    print(i)
    print('code')

将给予：

before
1
code
after
before
2
code
after
before
3
code
after

Python 仅进入和退出 with 块一次，因此您不能在进入/退出步骤中重复执行逻辑。

Answer 2

一个更完整的答案，如果异常可能发生在生成器之外：

from contextlib import contextmanager

class CustomError(RuntimeError):
    pass

@contextmanager
def handle_custom_error():
    try:
        yield
    except CustomError as e:
        print(f"handled: {e}")

def loop(seq):
    for i in seq:
        try:
            print('before')
            if i == 0:
                raise CustomError("inside generator")
            yield i # for body executes here
            print('after')
        except CustomError as e:
            print(f"handled: {e}")

@handle_custom_error()
def do_stuff(i):
    if i == 1:
        raise CustomError("inside do_stuff")
    print(f"i = {i}")

for i in loop(range(3)):
    do_stuff(i)

输出：

before
handled: inside generator
before
handled: inside do_stuff
after
before
i = 2
after

如何制作一个内部有循环的上下文管理器？

How do I make a contextmanager with a loop inside?

python

contextmanager