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 TstContx(object):
...    def __enter__(self):
...        raise Exception('Oops in __enter__')
...
...    def __exit__(self, e_typ, e_val, trcbak):
...        print "This isn't running"
... 
>>> with TstContx():
...     pass
... 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 3, in __enter__
Exception: Oops in __enter__
>>> 

Edit

This is as close as I could get...

class TstContx(object):
    def __enter__(self):
        try:
            # __enter__ code
        except Exception as e
            self.init_exc = e

        return self

    def __exit__(self, e_typ, e_val, trcbak):
        if all((e_typ, e_val, trcbak)):
            raise e_typ, e_val, trcbak

        # __exit__ code


with TstContx() as tc:
    if hasattr(tc, 'init_exc'): raise tc.init_exc

    # code in context

In hind sight, a context manager might have not been the best design decision

Cognizant

The problem is, it's not possible to skip with body from within __enter__ (see pep 377 ) – Ethe

Answer 1

Like this:

import sys

class Context(object):
    def __enter__(self):
        try:
            raise Exception("Oops in __enter__")
        except:
            # Swallow exception if __exit__ returns a True value
            if self.__exit__(*sys.exc_info()):
                pass
            else:
                raise


    def __exit__(self, e_typ, e_val, trcbak):
        print "Now it's running"


with Context():
    pass

To let the program continue on its merry way without executing the context block you need to inspect the context object inside the context block and only do the important stuff if __enter__ succeeded.

class Context(object):
    def __init__(self):
        self.enter_ok = True

    def __enter__(self):
        try:
            raise Exception("Oops in __enter__")
        except:
            if self.__exit__(*sys.exc_info()):
                self.enter_ok = False
            else:
                raise
        return self

    def __exit__(self, e_typ, e_val, trcbak):
        print "Now this runs twice"
        return True


with Context() as c:
    if c.enter_ok:
        print "Only runs if enter succeeded"

print "Execution continues"

As far as I can determine, you can't skip the with-block entirely. And note that this context now swallows all exceptions in it. If you wish not to swallow exceptions if __enter__ succeeds, check self.enter_ok in __exit__ and return False if it's True.

Epideictic

If there is an exception in the __enter__ and you call __exit__, is there any way to break out of the with block in the client code? – Cognizant

lol I just thought about that at the same time. I updated my question with the same logic. – Cognizant

If the goal is to execute the exit but not execute the with content, wouldn't you just do self.__exit__(*sys.exc_info()) and then raise the original exception regardless of the exit return value? Am I missing something? – Unaccompanied

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.

Dyanne

Answer 3

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

Edit: Unfortunately this approach still allows the user to create "dangling" resources that wont be cleaned up if he does something like:

foo = Foo() # this allocates resource without a with context.
raise ValueError("bla") # foo.__exit__() will never be called.

I am quite curious if this could be worked around by modifying the new implementation of the class or some other python magic that forbids object instantiation without a with context.

Photosynthesis

+1: This answer suited my use case the best, and makes a lot of sense. Surprised it does not have more upvotes. Thanks. – Breazeale

Well, technically I didn't answer the question ;-) – Photosynthesis

This looks rather a bit like an anti-pattern, because it enables using a resource without actually making use of a context manager. In many situations this can lead to writing code that fails to do the proper cleanup, which is the purpose of offering a context manager. The approach may be fine in cases where the execution of __enter__+__exit__ is indeed entirely optional. And admittedly the pattern is used frequently in the standard library, but that doesn't make it better ;). – Bondie

Yep, that is correct:) – Photosynthesis

Answer 4

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.

Busywork

Answer 5

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.

Unlimited

Answer 6

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

Gambell

yes, but this would throw "generator didn't yield" in contextlib. – Ethe

@thg435, reasonable, but we can wrap 'yield' with a try ... finally – Gambell

that is yield in a finally block – Gambell

there are plenty of workarounds, but the root of the problem is that it isn't possible to skip the entire with block. So even if we manage to handle exception in enter somehow, the block will still run, with None or some other rubbish as an argument. – Ethe

This is not working in Python 3.9 - would be nice to see update. – Alex

Answer 7

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.

Yellowish

Answer 8

There might be a better way to implement what you want.

By design, in the Python programming language,

When the __enter__ function raises an error, it means the resource to be acquired is not acquired, so there's no reason to call __exit__.

It's not explained in the question that in which case would you want to call __exit__, but the most common case is to make one context manager that handles two resources --- in other words, compose multiple context generators.

To which, the simplest solution in my opinion is to use @contextmanager and a function --- in fact, I don't even know how to write a class correctly.

@contextmanager
def nest_resource(a, b):
    with a as aa, b as bb:
        yield (aa, bb)

---

If the above is not the case for you, you can use the recipe Cleaning up in an __enter__ implementation in the documentation.

The important part is:

class ResourceManager:

    @contextmanager
    def _cleanup_on_error(self):
        with ExitStack() as stack:
            stack.push(self)
            yield
            stack.pop_all()

    def __enter__(self):
        resource = self.acquire_resource()
        with self._cleanup_on_error():
            if not self.check_resource_ok(resource):
                raise RuntimeError(msg.format(resource))
        return resource

    def __exit__(self, *exc_details):
        self.release_resource()

The way it works is:

• If an error happens in acquire_resource, then __exit__ will not be called.
• If an error happens in check_resource_ok (or that function returns False), then __exit__ will be called.
• If __enter__ returns peacefully, then __exit__ will not be called.

If you want to make __exit__ always be called whenever __enter__ fails, just put nothing before the with block.

Menzies