Sahand
Sahand

Reputation: 2175

Handling instances of a context manager inside another context manager

How should a context manager created inside another context manager be handled in Python?

Example: suppose you have class A that acts as a context manager, and class B that also acts as a context manager. But class B instances will have to instantiate and use an instance of class A. I've gone through PEP 343 and this is the solution I thought of:

class A(object):
    def __enter__(self):
        # Acquire some resources here
        return self

    def __exit__(seplf, exception_type, exception, traceback):
        # Release the resources and clean up
        pass


class B(object):
    def __init__(self):
        self.a = A()

    def __enter__(self):
        # Acquire some resources, but also need to "start" our instance of A
        self.a.__enter__()
        return self

    def __exit__(self, exception_type, exception, traceback):
        # Release the resources, and make our instance of A clean up as well
        self.a.__exit__(exception_type, exception, traceback)

Is this the correct approach? Or am I missing some gotchas?

Upvotes: 47

Views: 13877

Answers (4)

kuzzooroo
kuzzooroo

Reputation: 7408

If you can use the @contextlib.contextmanager decorator your life gets a lot easier:

import contextlib
    
@contextlib.contextmanager
def internal_cm():
    try:
        print("Entering internal_cm")
        yield None
        print("Exiting cleanly from internal_cm")
    finally:
        print("Finally internal_cm")

@contextlib.contextmanager
def external_cm():
    with internal_cm() as c:
        try:
            print("In external_cm_f", c)
            yield c
            print("Exiting cleanly from external_cm_f", c)
        finally:
            print("Finally external_cm_f", c) 

if "__main__" == __name__:
    with external_cm():
        print("Location A")
    print
    with external_cm():
        print("Location B")
        raise Exception("Some exception occurs!!")

Upvotes: 20

blim747
blim747

Reputation: 81

Expanding on kuzzooroo's top upvoted answer but applying it for Python 3+:

For convenience here is kuzzooroo's original code as Python 3 code (pretty much adding parenthesis for the print statements:

import contextlib

@contextlib.contextmanager
def internal_cm():
    try:
        print("Entering internal_cm")
        yield None
        print("Exiting cleanly from internal_cm")
    finally:
        print("Finally internal_cm")


@contextlib.contextmanager
def external_cm():
    with internal_cm() as c:
        try:
            print("In external_cm_f")
            yield [c]
            print("Exiting cleanly from external_cm_f")
        finally:
            print("Finally external_cm_f")

if "__main__" == __name__:
    with external_cm():
        print("Location A")
    with external_cm():
        print("Location B")
        raise Exception("Some exception occurs!!")

And here is the output of this script:

Entering internal_cm
In external_cm_f
Location A
Exiting cleanly from external_cm_f
Finally external_cm_f
Exiting cleanly from internal_cm
Finally internal_cm
Entering internal_cm
In external_cm_f
Location B
Finally external_cm_f
Finally internal_cm
Traceback (most recent call last):
  File "main.py", line 28, in <module>
    raise Exception("Some exception occurs!!")
Exception: Some exception occurs!!

Upvotes: 1

Noctis Skytower
Noctis Skytower

Reputation: 22041

Alternatively, you could write your code like so:

with A() as a:
    with B(a) as b:
        # your code here

Another solution you might want to try could be this:

class A:

    def __init__(self):
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

class B(A):

    def __init__(self):
        super().__init__()

    def __enter__(self):
        super().__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        super().__exit__(exc_type, exc_val, exc_tb)

After considering the explanation of your situation, this might be a better solution:

class Resource:

    def __init__(self, dependency=None):
        self.dependency = dependency
        # your code here

    def __enter__(self):
        if self.dependency:
            self.dependency.__enter__()
        # your code here
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # your code here
        if self.dependency:
            self.dependency.__exit__(exc_type, exc_val, exc_tb)

I am not sure if the following implementation is correct, but __exit__ must take proper care of exceptions. It is somewhat difficult for me to imagine how to recursively chain the calls while handling exceptions properly.

class Resource:

    def __init__(self, dependency=None):
        self.dependency = dependency
        self.my_init()

    def __enter__(self):
        if self.dependency:
            self.dependency.__enter__()
        return self.my_enter()

    def __exit__(self, exc_type, exc_val, exc_tb):
        suppress = False
        try:
            suppress = self.my_exit(exc_type, exc_val, exc_tb)
        except:
            exc_type, exc_val, exc_tb = sys.exc_info()
        if suppress:
            exc_type = exc_val = exc_tb = None
        if self.dependency:
            suppress = self.dependeny.__exit__(exc_type, exc_val, exc_tb)
            if not supress:
                raise exc_val.with_traceback(exc_tb) from None
        return suppress

    def my_init(self):
        pass

    def my_enter(self):
        pass

    def my_exit(self, exc_type, exc_val, exc_tb):
        pass

Upvotes: 10

user2394284
user2394284

Reputation: 6028

Here is an example of manual resource management in a contextmanager: The outer contextmanager manages the inner.

class Inner:

    def __enter__(self):
        print("<inner>")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print("</inner>")


class Outer:

    def __init__(self):
        self.inner = Inner()

    def __enter__(self):
        self.inner.__enter__()
        try:
            #raise RuntimeError("Suppose we fail here")
            print("<outer>")
            return self
        except Exception as e:
            self.inner.__exit__(None, None, None)
            raise e

    def __exit__(self, exc_type, exc_value, traceback):
        print("</outer>")
        self.inner.__exit__(exc_type, exc_value, traceback)

Usage is as normal:

with Outer() as scope:
    #raise RuntimeError("Suppose we fail here")
    pass

The observant reader will notice that the inner contextmanager now becomes a pointless marionette puppet (since we are pulling its threads manually). So be it.

Upvotes: 4

Related Questions