python - Correct way to manage multiple resources with context managers

Question

I have some resources which I have wrapped in a context manager 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}'

If I were to use 2 resources directly, I could use the following syntax:

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})')

This works as expected:

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

What I'd like to do, however, is wrap the use of these multiple resources into a class Task, and make that itself a context manager.

This is my first attempt at creating such a Task class which manages 2 resources:

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}')

I can now use the familiar syntax:

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

And again this works as expected:

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

This all works fine, except when an exception is thrown trying to free one of the resources.

To illustrate I have modified my Resource class to throw an exception for one of the resources:

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}'

With the previous manual use of 2 resources:

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

In the face of the exception freeing bar, foo is still freed:

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

However, doing the same with Task, I leak the second resource:

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

Output showing I never try free foo:

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

Questions:

• I could wrap each explicit call to Resource.__exit__ in a try/except block, but then if I want to propagate the exception further up the call stack whilst still freeing up all the resources, I have to keep track of all exceptions and then rethrow them afterwards... it feels incorrect.

• It also feels "dirty" calling Resource.__exit__ explicitly from Task.__exit__, rather than having the with statement call it implicitly for me. Is there a way to use a context manager inside a class like I'm trying to do?

What is the right way to handle multiple resources inside a single context manager?

标签： pythoncontextmanager

Answer 1

正如评论中提到的，ExitStack正是这样做的。

旨在使以编程方式组合其他上下文管理器变得容易的上下文管理器

您可以简单地继承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