Catching exception in context manager __enter__()

Question

Is it possible to ensure the __exit__() method is called even if there is an exception in __enter__()?

>>> class TstCont         


        

Answer 1

The docs contain an example that uses contextlib.ExitStack for ensuring the cleanup:

As noted in the documentation of ExitStack.push(), this method can be useful in cleaning up an already allocated resource if later steps in the __enter__() implementation fail.

So you would use ExitStack() as a wrapping context manager around the TstContx() context manager:

from contextlib import ExitStack

with ExitStack() as stack:
    ctx = TstContx()
    stack.push(ctx)  # Leaving `stack` now ensures that `ctx.__exit__` gets called.
    with ctx:
        stack.pop_all()  # Since `ctx.__enter__` didn't raise it can handle the cleanup itself.
        ...  # Here goes the body of the actual context manager.

Answer 2

No. If there is the chance that an exception could occur in __enter__() then you will need to catch it yourself and call a helper function that contains the cleanup code.

Answer 3

You could use contextlib.ExitStack (not tested):

with ExitStack() as stack:
    cm = TstContx()
    stack.push(cm) # ensure __exit__ is called
    with ctx:
         stack.pop_all() # __enter__ succeeded, don't call __exit__ callback

Or an example from the docs:

stack = ExitStack()
try:
    x = stack.enter_context(cm)
except Exception:
    # handle __enter__ exception
else:
    with stack:
        # Handle normal case

See contextlib2 on Python <3.3.

Answer 4

if inheritance or complex subroutines are not required, you can use a shorter way:

from contextlib import contextmanager

@contextmanager
def test_cm():
    try:
        # dangerous code
        yield  
    except Exception, err
        pass # do something

Answer 5

class MyContext:
    def __enter__(self):
        try:
            pass
            # exception-raising code
        except Exception as e:
            self.__exit__(e)

    def __exit__(self, *args):
        # clean up code ...
        if args[0]:
            raise

I've done it like this. It calls __exit__() with the error as the argument. If args[0] contains an error it reraises the exception after executing the clean up code.

Answer 6

I suggest you follow RAII (resource acquisition is initialization) and use the constructor of your context to do the potentially failing allocation. Then your __enter__ can simply return self which should never ever raise an exception. If your constructor fails, the exception may be thrown before even entering the with context.

class Foo:
    def __init__(self):
        print("init")
        raise Exception("booh")

    def __enter__(self):
        print("enter")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("exit")
        return False


with Foo() as f:
    print("within with")

Output:

init
Traceback (most recent call last):
  File "<input>", line 1, in <module>
  ...
    raise Exception("booh")
Exception: booh