Question

I have little experience with decorators in Python, but I'd like to write a function decorator that runs the function, catches a specific exception, and if the exception is caught then re-tries the function a certain number of times. That is, I'd like to do this:

@retry_if_exception(BadStatusLine, max_retries=2)
def thing_that_sometimes_fails(self, foo):
   foo.do_something_that_sometimes_raises_BadStatusLine()

I assume this kind of thing is easy with decorators, but I'm not clear about how exactly to go about it.

Was it helpful?

Solution

from functools import wraps
def retry_if_exception(ex, max_retries):
    def outer(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            assert max_retries > 0
            x = max_retries
            while x:
                try:
                    return func(*args, **kwargs)
                except ex:
                    x -= 1
        return wrapper
    return outer

see why you better use @wraps

OTHER TIPS

I think you're basically wanting something like this:

def retry_if_exception(exception_type=Exception, max_retries=1):
    def decorator(fn):
        def wrapper(*args, **kwargs):
            for i in range(max_retries+1):
                print('Try #', i+1)
                try:
                    return fn(*args, **kwargs)
                except exception_type as e:
                    print('wrapper exception:', i+1, e)
        return wrapper
    return decorator

@retry_if_exception()
def foo1():
    raise Exception('foo1')

@retry_if_exception(ArithmeticError)
def foo2():
    x=1/0

@retry_if_exception(Exception, 2)
def foo3():
    raise Exception('foo3')

The following seems to do what you've described:

def retry_if_exception( exception, max_retries=2 ):
    def _retry_if_exception( method_fn ):
        # method_fn is the function that gives rise
        # to the method that you've decorated,
        # with signature (slf, foo)
        from functools import wraps
        def method_deco( slf, foo ):
            tries = 0
            while True:
                try:
                    return method_fn(slf, foo)
                except exception:
                    tries += 1
                    if tries > max_retries:
                        raise
        return wraps(method_fn)(method_deco)
    return _retry_if_exception

Here's an example of it in use:

d = {}

class Foo():
    def usually_raise_KeyError(self):
        print("d[17] = %s" % d[17])

foo1 = Foo()

class A():
    @retry_if_exception(KeyError, max_retries=2)
    def something_that_sometimes_fails( self, foo ):
        print("About to call foo.usually_raise_KeyError()")
        foo.usually_raise_KeyError()

a = A()
a.something_that_sometimes_fails(foo1)

This gives:

About to call foo.usually_raise_KeyError()
About to call foo.usually_raise_KeyError()
About to call foo.usually_raise_KeyError()
Traceback (most recent call last):
  File " ......... TrapRetryDeco.py", line 39, in <module>
    a.something_that_sometimes_fails( foo1)
  File " ......... TrapRetryDeco.py", line 15, in method_deco
    return method_fn( slf, foo)
  File " ......... TrapRetryDeco.py", line 36, in something_that_sometimes_fails
    foo.usually_raise_KeyError()
  File " ......... TrapRetryDeco.py", line 28, in usually_raise_KeyError
    print("d[17] = %s" % d[17])
KeyError: 17

I assume that by "2 retries" you mean the operation will be attempted 3x all told. Your example has a couple of complications which may obscure the basic setup: It seems you want a method decorator, as your function/method's first parameter is "self"; however, that method immediately delegates to some bad method of its foo parameter. I preserved these complications :)

As outline, you would do something along these lines:

import random

def shaky():
    1/random.randint(0,1)

def retry_if_exception(f):
    def inner(retries=2):
        for retry in range(retries):
            try:
                return f()
            except ZeroDivisionError:
                print 'try {}'.format(retry)
        raise         

    return inner            

@retry_if_exception
def thing_that_may_fail():
    shaky()

thing_that_may_fail() 

As written, that will fail about 1/2 the time.

When it does fail, prints:

try 0
try 1
Traceback (most recent call last):
  File "Untitled 2.py", line 23, in <module>
    thing_that_may_fail()    
  File "Untitled 2.py", line 10, in inner
    return f()
  File "Untitled 2.py", line 21, in thing_that_may_fail
    shaky()
  File "Untitled 2.py", line 4, in shaky
    1/random.randint(0,1)
ZeroDivisionError: integer division or modulo by zero

You could adapt this structure to many different types of errors.

Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top