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 ?

Do you want the body of the `with` statement to execute in its entirety if an exception is raised in the middle, or simply catch the exception and continue after the `with` statement? — chepner, Feb 23 '20 at 22:44
I know I am too late but this might help: https://book.pythontips.com/en/latest/context_managers.html#handling-exceptions — Ganesh Patil, Oct 06 '22 at 11:33

kaya3 · Accepted Answer · 2020-02-23T22:46:13.087

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.

Also want to add that if you don't have a resource to `yield`, then an empty yield statement works as well — smac89, Jun 30 '22 at 19:22
I'd never seen `else:` after a `try:` block -- If you're more used to exception handling in i.e. java, could just put the "cleanup that should only be done on success" code on the line after `yield` — Carl Walsh, Jun 06 '23 at 22:37
@CarlWalsh Yes; for clarity, the use of `else` in Python is that it happens after the `try` and only on success, but any exceptions thrown in the `else` block won't be caught by the preceding `except` block. So an exception raised during cleanup won't be treated the same as an exception raised in the external `with` block. — kaya3, Jun 06 '23 at 23:39

tdelaney · Answer 2 · 2020-02-24T01:10:48.393

Instead of using contextmanager, which I find a bit clunky for this task (you need to fake a generator with a yield just to log an exception), write a context manager yourself. Its simply a class with two special methods. And with the predefined AbstractContextManager you only need to implement one of them:

import contextlib

class ExceptionLogger(contextlib.AbstractContextManager):

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            print("***Logging exception {}***".format((exc_type, exc_value, 
                traceback)))

with ExceptionLogger():
    raise ValueError("foo")

Context Manager that handles exceptions

2 Answers2

Linked