Merge pull request #127 from Brian-Williams/master

Add retry_if_exception_message

tenacity/__init__.py

@@ -47,6 +47,8 @@
 from .retry import retry_if_result  # noqa
 from .retry import retry_never  # noqa
 from .retry import retry_unless_exception_type  # noqa
+from .retry import retry_if_exception_message  # noqa
+from .retry import retry_if_not_exception_message  # noqa
 
 # Import all nap strategies for easier usage.
 from .nap import sleep  # noqa

tenacity/retry.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import abc
+import re
 
 import six
 
@@ -111,6 +112,50 @@ def __call__(self, attempt):
             return not self.predicate(attempt.result())
 
 
+class retry_if_exception_message(retry_if_exception):
+    """Retries if an exception message equals or matches."""
+
+    def __init__(self, message=None, match=None):
+        if message and match:
+            raise TypeError(
+                "{}() takes either 'message' or 'match', not both".format(
+                    self.__class__.__name__))
+
+        # set predicate
+        if message:
+            def message_fnc(exception):
+                return message == str(exception)
+            predicate = message_fnc
+        elif match:
+            prog = re.compile(match)
+
+            def match_fnc(exception):
+                return prog.match(str(exception))
+            predicate = match_fnc
+        else:
+            raise TypeError(
+                "{}() missing 1 required argument 'message' or 'match'".
+                format(self.__class__.__name__))
+
+        super(retry_if_exception_message, self).__init__(predicate)
+
+
+class retry_if_not_exception_message(retry_if_exception_message):
+    """Retries until an exception message equals or matches."""
+
+    def __init__(self, *args, **kwargs):
+        super(retry_if_not_exception_message, self).__init__(*args, **kwargs)
+        # invert predicate
+        if_predicate = self.predicate
+        self.predicate = lambda *args_, **kwargs_: not if_predicate(
+            *args_, **kwargs_)
+
+    def __call__(self, attempt):
+        if not attempt.failed:
+            return True
+        return self.predicate(attempt.exception())
+
+
 class retry_any(retry_base):
     """Retries if any of the retries condition is valid."""
 

tenacity/tests/test_tenacity.py

@@ -473,6 +473,15 @@ def _r():
                           _r)
         self.assertEqual(5, r.statistics['attempt_number'])
 
+    def test_retry_if_exception_message_negative_no_inputs(self):
+        with self.assertRaises(TypeError):
+            tenacity.retry_if_exception_message()
+
+    def test_retry_if_exception_message_negative_too_many_inputs(self):
+        with self.assertRaises(TypeError):
+            tenacity.retry_if_exception_message(
+                message="negative", match="negative")
+
 
 class NoneReturnUntilAfterCount(object):
     """Holds counter state for invoking a method several times in a row."""
@@ -531,6 +540,8 @@ def go(self):
 class NameErrorUntilCount(object):
     """Holds counter state for invoking a method several times in a row."""
 
+    derived_message = "Hi there, I'm a NameError"
+
     def __init__(self, count):
         self.counter = 0
         self.count = count
@@ -543,7 +554,7 @@ def go(self):
         if self.counter < self.count:
             self.counter += 1
             return True
-        raise NameError("Hi there, I'm a NameError")
+        raise NameError(self.derived_message)
 
 
 class IOErrorUntilCount(object):
@@ -579,12 +590,14 @@ def __init__(self, value):
         self.value = value
 
     def __str__(self):
-        return repr(self.value)
+        return self.value
 
 
 class NoCustomErrorAfterCount(object):
     """Holds counter state for invoking a method several times in a row."""
 
+    derived_message = "This is a Custom exception class"
+
     def __init__(self, count):
         self.counter = 0
         self.count = count
@@ -596,8 +609,7 @@ def go(self):
         """
         if self.counter < self.count:
             self.counter += 1
-            derived_message = "This is a Custom exception class"
-            raise CustomError(derived_message)
+            raise CustomError(self.derived_message)
         return True
 
 
@@ -657,6 +669,38 @@ def _retryable_test_with_unless_exception_type_no_input(thing):
     return thing.go()
 
 
+@retry(
+    stop=tenacity.stop_after_attempt(5),
+    retry=tenacity.retry_if_exception_message(
+        message=NoCustomErrorAfterCount.derived_message))
+def _retryable_test_if_exception_message_message(thing):
+    return thing.go()
+
+
+@retry(retry=tenacity.retry_if_not_exception_message(
+    message=NoCustomErrorAfterCount.derived_message))
+def _retryable_test_if_not_exception_message_message(thing):
+    return thing.go()
+
+
+@retry(retry=tenacity.retry_if_exception_message(
+    match=NoCustomErrorAfterCount.derived_message[:3] + ".*"))
+def _retryable_test_if_exception_message_match(thing):
+    return thing.go()
+
+
+@retry(retry=tenacity.retry_if_not_exception_message(
+    match=NoCustomErrorAfterCount.derived_message[:3] + ".*"))
+def _retryable_test_if_not_exception_message_match(thing):
+    return thing.go()
+
+
+@retry(retry=tenacity.retry_if_not_exception_message(
+    message=NameErrorUntilCount.derived_message))
+def _retryable_test_not_exception_message_delay(thing):
+    return thing.go()
+
+
 @retry
 def _retryable_default(thing):
     return thing.go()
@@ -773,6 +817,49 @@ def test_retry_until_exception_of_type_wrong_exception(self):
             self.assertTrue(isinstance(e, RetryError))
             print(e)
 
+    def test_retry_if_exception_message(self):
+        try:
+            self.assertTrue(_retryable_test_if_exception_message_message(
+                NoCustomErrorAfterCount(3)))
+        except CustomError:
+            print(_retryable_test_if_exception_message_message.retry.
+                  statistics)
+            self.fail("CustomError should've been retried from errormessage")
+
+    def test_retry_if_not_exception_message(self):
+        try:
+            self.assertTrue(_retryable_test_if_not_exception_message_message(
+                NoCustomErrorAfterCount(2)))
+        except CustomError:
+            s = _retryable_test_if_not_exception_message_message.retry.\
+                statistics
+            self.assertTrue(s['attempt_number'] == 1)
+
+    def test_retry_if_not_exception_message_delay(self):
+        try:
+            self.assertTrue(_retryable_test_not_exception_message_delay(
+                NameErrorUntilCount(3)))
+        except NameError:
+            s = _retryable_test_not_exception_message_delay.retry.statistics
+            print(s['attempt_number'])
+            self.assertTrue(s['attempt_number'] == 4)
+
+    def test_retry_if_exception_message_match(self):
+        try:
+            self.assertTrue(_retryable_test_if_exception_message_match(
+                NoCustomErrorAfterCount(3)))
+        except CustomError:
+            self.fail("CustomError should've been retried from errormessage")
+
+    def test_retry_if_not_exception_message_match(self):
+        try:
+            self.assertTrue(_retryable_test_if_not_exception_message_message(
+                NoCustomErrorAfterCount(2)))
+        except CustomError:
+            s = _retryable_test_if_not_exception_message_message.retry.\
+                statistics
+            self.assertTrue(s['attempt_number'] == 1)
+
     def test_defaults(self):
         self.assertTrue(_retryable_default(NoNameErrorAfterCount(5)))
         self.assertTrue(_retryable_default_f(NoNameErrorAfterCount(5)))