Context Manager that handles exceptions

Question

I am trying to wrap my head around how to write a context manager that deals with writing some logs while handling any exceptions. The problem I am trying to solve is to make code like this:

try:
    # code that can raise exception here
except Exception as e:
    print('failed', e)

print('all good')

This is a repeated pattern I have in the code and I think it's best handled with a context manager like:

with my_ctx_manager(success_msg='all good', failed_msg='failed):
    # code that can raise exception here

this looks much better, but I don't know how to write the actual context manager to deal with any exceptions that could rise inside the context.

@contextlib.contextmanager
def my_ctx_manager(success_msg, failed_msg):
   try:
       # if no exception then print(success_msg)
       # How do I catch any exception here
   except Exception:
      print(failed_msg)
      # I need the exception to propagate as well
      raise

I guess my question is more of the type: How do I make sure that the context manager correctly catches, logs and re-raise any exception for the code that is wrapping ?

Answer 1

The way the @contextmanager decorator works, you should write yield once within your context manager function, so that the with block will be executed while the yield statement pauses your function's execution. That means if the with block throws an exception, you can catch it by wrapping yield in a try/except block:

from contextlib import contextmanager

@contextmanager
def example():
    print('entered the context manager')
    managed_resource = 'some resource'
    try:
        yield managed_resource
    except Exception as e:
        print('caught:', e)
        # any cleanup that should only be done on failure
        raise
    else:
        # any cleanup that should only be done on success
        print('no exception was thrown')
    finally:
        # any cleanup that should always be done
        print('exited the context manager')

with example() as resource:
    print('resource:', resource)
    raise ValueError('some error message')

Output:

entered the context manager
resource: some resource
caught: some error message
exited the context manager
Traceback (most recent call last):
  File "<stdin>", line 3, in <module>
ValueError: some error message

If you want to catch everything (not just Exception), then you can write a bare except: block and use sys.exc_info() to get the exception information.