Search code examples
pythonexceptioncontextmanager

Changing the exception type in a context manager


I currently have code that looks like this scattered throughout my codebase:

try: 
    something()
except Exception as exc:
    raise SomethingError from exc

I would like to write a context manager that would remove some of this boiler-plate:

with ExceptionWrapper(SomethingError):
    something()

It looks like it is possible to suppress exceptions inside a context manager - see: contextlib.suprress. It doesn't look like it is possible to change what exception is being raised.

However, I haven't been able to find clear documentation on what the return value of the __exit__ function of a context manager is.


Solution

  • This sort of simple context manager is easiest to implement using contextlib.contextmanager, which creates a context manager out of a generator. Simply wrap the first yield statement in the generator with a try block and raise the desired exception in the except block:

    import contextlib
    from typing import Iterator
    
    
    class SomethingError(Exception):
        pass
    
    
    @contextlib.contextmanager
    def exception_wrapper(message: str) -> Iterator[None]:
        """Wrap Exceptions with SomethingError."""
        try: 
            yield
        except Exception as exc:
            raise SomethingError(message) from exc
    

    This can be used like so:

    >>> with exception_wrapper("and eggs!"):
    ...     raise Exception("spam")
    
    Traceback (most recent call last):
      File "/path/to/script.py", line 11, in exception_wrapper
        yield
      File "/path/to/script.py", line 17, in <module>
        raise Exception("spam")
    Exception: spam
    
    The above exception was the direct cause of the following exception:
    
    Traceback (most recent call last):
      File "/path/to/script.py", line 16, in <module>
        with exception_wrapper("and eggs!"):
      File "/usr/lib/python3.10/contextlib.py", line 153, in __exit__
        self.gen.throw(typ, value, traceback)
      File "/path/to/script.py", line 13, in exception_wrapper
        raise SomethingError(message) from exc
    __main__.SomethingError: and eggs!