如何创建一个包含循环的上下文管理器？

Question

如何创建一个包含循环的上下文管理器？

15

我想要类似这样的东西：

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。需要一种方法对一些任意代码进行循环，并对每个迭代进行设置/清理和异常处理，且与以上with语句一样简单。

- jpkotta

2个回答

5

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

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

- jpkotta

这是一个很好的解决方案。然而，它需要将 do_stuff 代码放在单独的修饰方法中。如何使用更简单的方法来处理这个问题呢？ for i in loop(range(3)): if i == 1: raise CustomError("inside code") 是否有一种方法可以在Python中处理此问题，而无需在用户代码中同时使用 for 和嵌套 with？（我认为 ruby 的代码块会允许这样做，因此我提出了这个问题） - Mirek

基本上我正在尝试增强我的测试框架，以便QA开发人员可以轻松地使用一个简单的结构来重试可能会引发AssertionError（或其他自定义异常元组）的代码，因此生成器和上下文管理器必须共享一些范围以了解要处理哪个异常。 - Mirek

@Mirek 我不知道是否有办法做到这一点。实际上，这正是我想要做的，但我的答案已经足够满足我的需求了（我想要捕获的所有错误都在预定义的函数中，而且它们在这些函数之外几乎不可能发生）。如果你想要做你想要的事情，你可以先看看 fuckit.py 是如何工作的。 - jpkotta

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Simeon Visser · Accepted Answer

我认为在这里使用发电机效果更好：

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 块一次，因此您不能在进入/退出步骤中拥有会重复执行的逻辑。