使用上下文管理器管理多个资源的正确方法

Question

我有一些资源包装在上下文管理器中 class。

class Resource:
    def __init__(self, res):
        print(f'allocating resource {res}')
        self.res = res

    def __enter__(self):
        return self.res

    def __exit__(self, typ, value, traceback):
        print(f'freed resource {self.res}')

    def __str__(self):
        return f'{self.res}'

如果我要直接使用2个资源，我可以使用以下语法：

with Resource('foo') as a, Resource('bar') as b:
    print(f'doing something with resource a({a})')
    print(f'doing something with resource b({b})')

这按预期工作：

allocating resource foo
allocating resource bar
doing something with resource a(foo)
doing something with resource b(bar)
freed resource bar
freed resource foo

然而，我想做的是将这些多种资源的使用包装到 class Task 中，并使其本身成为上下文管理器。

这是我第一次尝试创建这样一个 Task class 来管理 2 个资源：

class Task:
    def __init__(self, res1, res2):
        self.a = Resource(res1)
        self.b = Resource(res2)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.b.__exit__(type, value, traceback)
        self.a.__exit__(type, value, traceback)

    def run(self):
        print(f'running task with resource {self.a} and {self.b}')

我现在可以使用熟悉的语法了：

with Task('foo', 'bar') as t:
    t.run()

再次按预期工作：

allocating resource foo
allocating resource bar
running task with resource foo and bar
freed resource bar
freed resource foo

一切正常，除非在尝试释放其中一个资源时抛出异常。

为了说明，我修改了我的 Resource class 以针对以下资源之一抛出异常：

class Resource:
    def __init__(self, res):
        print(f'allocating resource {res}')
        self.res = res

    def __enter__(self):
        return self.res

    def __exit__(self, typ, value, traceback):
        print(f'try free resource {self.res}')
        if self.res == 'bar':
            raise RuntimeError(f'error freeing {self.res} resource')
        print(f'freed resource {self.res}')

    def __str__(self):
        return f'{self.res}'

加上之前手动使用2个资源：

try:
    with Resource('foo') as a, Resource('bar') as b:
        print(f'doing something with resource a({a})')
        print(f'doing something with resource b({b})')
except:
    pass

面对释放bar的异常，foo仍然被释放：

allocating resource foo
allocating resource bar
doing something with resource a(foo)
doing something with resource b(bar)
try free resource bar
try free resource foo
freed resource foo

然而，对Task做同样的事情，我泄露了第二个资源：

try:
    with Task('foo', 'bar') as t:
        t.run()
except:
    pass

输出显示我从不免费试用 foo:

allocating resource foo
allocating resource bar
running task with resource foo and bar
try free resource bar

问题：

我可以将对 Resource.__exit__ 的每个显式调用包装在 try/except 块中，但是如果我想进一步传播异常，调用堆栈，同时仍然释放所有资源，我必须跟踪所有异常，然后在之后重新抛出它们……感觉不正确。
也感觉 "dirty" 从 Task.__exit__ 显式调用 Resource.__exit__，而不是让 with 语句为我隐式调用它。有没有办法像我想做的那样在 inside a class 中使用上下文管理器？

在单个上下文管理器中处理多个资源的正确方法是什么？

Answer 1

如评论中所述，ExitStack 正是这样做的。

A context manager that is designed to make it easy to programmatically combine other context managers

您可以简单地从 ExitStack 继承并为您想要管理的每个资源调用 enter_context：

class Task(contextlib.ExitStack):
    def __init__(self, res1, res2):
        super().__init__()
        self.a = self.enter_context(Resource(res1))
        self.b = self.enter_context(Resource(res2))

    def run(self):
        print(f'running task with resource {self.a} and {self.b}')

请注意，无需定义您自己的 __enter__ 和 __exit__ 函数，因为 ExitStack 已为我们完成。

在示例中使用它：

try:
    with Task('foo', 'bar') as t:
        t.run()
except:
    pass

现在抛出释放 bar 的异常时，foo 仍然被释放：

allocating resource foo
allocating resource bar
running task with resource foo and bar
try free resource bar
try free resource foo
freed resource foo

使用上下文管理器管理多个资源的正确方法

Correct way to manage multiple resources with context managers

python

contextmanager