Handling exceptions inside context managers

Mauro Baraldi picture Mauro Baraldi · Feb 18, 2016 · Viewed 24.1k times · Source

I have some code where I try to reach a resource but sometimes it is unavailable, and results in an exception. I tried to implement a retry engine using context managers, but I can't handle the exception raised by the caller inside the __enter__ context for my context manager.

class retry(object):
    def __init__(self, retries=0):
        self.retries = retries
        self.attempts = 0
    def __enter__(self):
        for _ in range(self.retries):
            try:
                self.attempts += 1
                return self
            except Exception as e:
                err = e
    def __exit__(self, exc_type, exc_val, traceback):
        print 'Attempts', self.attempts

These are some examples which just raise an exception (which I expected to handle)

>>> with retry(retries=3):
...     print ok
... 
Attempts 1
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
NameError: name 'ok' is not defined
>>> 
>>> with retry(retries=3):
...     open('/file')
... 
Attempts 1
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
IOError: [Errno 2] No such file or directory: '/file'

Is there any way to intercept these exception(s) and handle them inside the context manager?

Answer

thefourtheye picture thefourtheye · Feb 18, 2016

Quoting __exit__,

If an exception is supplied, and the method wishes to suppress the exception (i.e., prevent it from being propagated), it should return a true value. Otherwise, the exception will be processed normally upon exit from this method.

By default, if you don't return a value explicitly from a function, Python will return None, which is a falsy value. In your case, __exit__ returns None and that is why the exeception is allowed to flow past the __exit__.

So, return a truthy value, like this

class retry(object):

    def __init__(self, retries=0):
        ...


    def __enter__(self):
        ...

    def __exit__(self, exc_type, exc_val, traceback):
        print 'Attempts', self.attempts
        print exc_type, exc_val
        return True                                   # or any truthy value

with retry(retries=3):
    print ok

the output will be

Attempts 1
<type 'exceptions.NameError'> name 'ok' is not defined

If you want to have the retry functionality, you can implement that with a generator, like this

def retry(retries=3):
    left = {'retries': retries}

    def decorator(f):
        def inner(*args, **kwargs):
            while left['retries']:
                try:
                    return f(*args, **kwargs)
                except NameError as e:
                    print e
                    left['retries'] -= 1
                    print "Retries Left", left['retries']
            raise Exception("Retried {} times".format(retries))
        return inner
    return decorator


@retry(retries=3)
def func():
    print ok

func()