Merge pull request #37 from harlowja/allow-retry-on-exception-exception-tuple

Allow providing exception tuple for 'retry_on_exception' argument

2 files changed, 15 insertions, 4 deletions

diff --git a/retrying.py b/retrying.py
index 594bf71..bcb7a9d 100644
--- a/retrying.py
+++ b/retrying.py
@@ -23,6 +23,12 @@ import traceback
 MAX_WAIT = 1073741823
 
 
+def _retry_if_exception_of_type(retryable_types):
+    def _retry_if_exception_these_types(exception):
+        return isinstance(exception, retryable_types)
+    return _retry_if_exception_these_types
+
+
 def retry(*dargs, **dkw):
     """
     Decorator function that instantiates the Retrying object
@@ -133,9 +139,14 @@ class Retrying(object):
         if retry_on_exception is None:
             self._retry_on_exception = self.always_reject
         else:
+            # this allows for providing a tuple of exception types that
+            # should be allowed to retry on, and avoids having to create
+            # a callback that does the same thing
+            if isinstance(retry_on_exception, (tuple)):
+                retry_on_exception = _retry_if_exception_of_type(
+                    retry_on_exception)
             self._retry_on_exception = retry_on_exception
 
-        # TODO simplify retrying by Exception types
         # retry on result filter
         if retry_on_result is None:
             self._retry_on_result = self.never_reject
diff --git a/test_retrying.py b/test_retrying.py
index 3cc843a..8ce4ac3 100644
--- a/test_retrying.py
+++ b/test_retrying.py
@@ -252,7 +252,7 @@ def _retryable_test_with_stop(thing):
     return thing.go()
 
 
-@retry(retry_on_exception=retry_if_exception_of_type(IOError))
+@retry(retry_on_exception=(IOError,))
 def _retryable_test_with_exception_type_io(thing):
     return thing.go()
 
@@ -264,14 +264,14 @@ def _retryable_test_with_exception_type_io_wrap(thing):
 
 @retry(
     stop_max_attempt_number=3,
-    retry_on_exception=retry_if_exception_of_type(IOError))
+    retry_on_exception=(IOError,))
 def _retryable_test_with_exception_type_io_attempt_limit(thing):
     return thing.go()
 
 
 @retry(
     stop_max_attempt_number=3,
-    retry_on_exception=retry_if_exception_of_type(IOError),
+    retry_on_exception=(IOError,),
     wrap_exception=True)
 def _retryable_test_with_exception_type_io_attempt_limit_wrap(thing):
     return thing.go()