add support for retry

decorators/__init__.py

@@ -1,3 +1,3 @@
 from .decorators import (
-    timeout_decorator, db_commit_decorator, parse_decorator, timeout
+    timeout_decorator, db_commit_decorator, parse_decorator, timeout, retry
 )

decorators/decorators.py

@@ -1,9 +1,11 @@
-from functools import wraps
+import time
+import collections
+from functools import wraps, partial
 from traceback import format_tb
 
 from db.basic import db_session
 from logger import (
-    parser, crawler, storage)
+    parser, crawler, storage, other)
 from utils import KThread
 from exceptions import Timeout
 
@@ -90,3 +92,32 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return crwal_decorator
+
+
+def retry(times=-1, delay=0, exceptions=Exception, logger=other):
+    """
+    inspired by https://github.com/invl/retry
+    :param times: retry times
+    :param delay: internals between each retry
+    :param exceptions: exceptions may raise in retry
+    :param logger: log for retry
+    :return: func result or None
+    """
+    def _inter_retry(caller, retry_time, retry_delay, es):
+        while retry_time:
+            try:
+                return caller()
+            except es as e:
+                retry_time -= 1
+                if not retry_time:
+                    logger.error("max tries for {} times, {} is raised, details: func name is {}, func args are {}".
+                                 format(times, e, caller.func.__name__, (caller.args, caller.keywords)))
+                    raise
+                time.sleep(retry_delay)
+
+    def retry_oper(func):
+        @wraps(func)
+        def _wraps(*args, **kwargs):
+            return _inter_retry(partial(func, *args, **kwargs), times, delay, exceptions)
+        return _wraps
+    return retry_oper

tests/test_utils.py

@@ -1,9 +1,15 @@
 import os
 import time
 
+try:
+    from unittest import mock
+except ImportError:
+    import mock
 import pytest
 
-from decorators import timeout
+from exceptions import CookieGenException
+from decorators import (
+    timeout, retry)
 from utils import (
     send_email, url_filter, text_filter)
 
@@ -59,3 +65,40 @@ def test_url_filter(url, expect):
 )
 def test_text_filter(text, expect):
     assert text_filter(text) == expect
+
+
+def test_retry_decorator():
+    hit = [0]
+    tries = 5
+    i = 1
+    j = 0
+
+    @retry(tries, exceptions=ZeroDivisionError)
+    def f(a, b):
+        hit[0] += 1
+        return a / b
+
+    with pytest.raises(ZeroDivisionError):
+        f(i, j)
+
+    assert hit[0] == tries
+
+
+def test_retry_with_delay():
+    get_cookies = mock.Mock()
+    get_cookies.side_effect = CookieGenException('Failed to gen cookies')
+
+    tries = 5
+    delay = 2
+    total_sleep_time = [0]
+
+    @retry(tries, delay, exceptions=CookieGenException)
+    def get():
+        total_sleep_time[0] += delay
+        return get_cookies()
+
+    with pytest.raises(CookieGenException):
+        get()
+
+    assert total_sleep_time[0] == tries * delay
+

utils/util_cls.py

@@ -3,54 +3,37 @@
 
 
 class KThread(threading.Thread):
-
     def __init__(self, *args, **kwargs):
-
         threading.Thread.__init__(self, *args, **kwargs)
-
         self.killed = False
 
     def start(self):
-
         """Start the thread."""
-
         self.__run_backup = self.run
-
         self.run = self.__run  # Force the Thread to install our trace.
 
         threading.Thread.start(self)
 
     def __run(self):
 
-        """Hacked run function, which installs the
-
-        trace."""
-
+        """Hacked run function, which installs the trace."""
         sys.settrace(self.globaltrace)
-
         self.__run_backup()
-
         self.run = self.__run_backup
 
     def globaltrace(self, frame, why, arg):
-
         if why == 'call':
-
             return self.localtrace
 
         else:
-
             return None
 
     def localtrace(self, frame, why, arg):
-
         if self.killed:
-
             if why == 'line':
                 raise SystemExit()
 
         return self.localtrace
 
     def kill(self):
-
         self.killed = True