Getting a context manager created with @contextmanager to work properly with exceptions

Question

I have the following code

from contextlib import contextmanager

@contextmanager
def simple_context_manager():
    print("starting context manager")
    yield
    print("finished context manager")

try:
    with simple_context_manager():
        raise RuntimeError
except RuntimeError:
    print("Caught the error")
print("Moving on")

Right now it prints out

starting context manager
Caught the error
Moving on

which tells me that the context manager isn't closing. How can I get it to close and print the "finished context manager" line?

Since I'm using the decorator, I don't have a dedicated __exit__ function that I think should be called according to this.

So I'm not sure how to get my context manager to exit in the case of an error occurring inside of its context.

Answer 1

You need a try-finally:

@contextmanager
def simple_context_manager():
    print("starting context manager")
    try:
        yield
    finally:
        print("finished context manager")

If an exception propagates out of the with statement, the @contextmanager decorator will throw the exception into the decorated generator at the point of the yield. finally lets us execute a cleanup block regardless of whether an exception occurs, so we use that.