python3 context manager force early exit

Question

I need to create a context manager that, when certain conditions are met, can be forced to exit early.

More details:

The context manager needs to handle checking/locking/releasing a resource. On __enter__, the context manager needs to check if the resource is locked. If it is, I'd like to have __exit__ called without executing the code in the context. Otherwise, the context manager acquires the resource, executes the context code, and cleans up the resource in __exit__.

It might look something like this:

class my_context_manager:
    def __enter__(self):
        if resource_locked():
            self.__exit__(None, ResourceLockedException(), None)
        else:
            acquire_resource()

    def __exit__(self, *exc_info):
        if not isinstance(exc_info[1], ResourceLockedException):
            release_resource()
        else:
            log.warn("Resource already in use.")

The code above doesn't actually work, however, because calling __exit__ inside of __enter__ doesn't stop the code within the context from being executed.

Alternatively, I could throw ResourceLockedException from within __enter__, but then __exit__ won't be called, since the exception would be thrown from the context manager itself. I'd like to be able to catch the exception, log a warning, and not enter the context if the resource is locked.

This comes down to finding some way of closing the context early, so that __exit__ is called and the context code isn't executed. Is there some way to tweak either of the ideas above to do this? Or is there another way?

Answer 1

Yes, you can't manually call __exit__ like this. One other alternative is simply raise the exception and let another construct manage it. You could either use a try-except or whip up another context manager to log these:

from contextlib import contextmanager

@contextmanager
def log_exception():
    try:
        yield
    except ResourceLockedException as e:
        log.warn(e)

and change your original context manager to:

class my_context_manager:
    def __enter__(self):
        if True:
            raise ResourceLockedException("Resource already in use")
        acquire_resource()
    def __exit__(self):
        release_resource()

And call it with:

with log_exception(), my_context_manager():
    # do things when resource acquired. 

Of course you could simply use a try-except and nest with inside that or use an if clause alternatively.