diff --git a/hexlib/misc.py b/hexlib/misc.py index 29667e7..fd7981a 100644 --- a/hexlib/misc.py +++ b/hexlib/misc.py @@ -11,13 +11,16 @@ last_time_called = dict() def retry(attempts, callback=None): def decorate(func): - retries = attempts - while retries > 0: - try: - func() - except Exception as e: - if callback: - callback(e) + def wrapper(*args, **kwargs): + retries = attempts + while retries > 0: + try: + return func(*args, **kwargs) + except Exception as e: + if callback: + callback(e) + retries -= 1 + return wrapper return decorate diff --git a/setup.py b/setup.py index 3a1a1ce..de9322b 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup setup( name="hexlib", - version="1.22", + version="1.23", description="Misc utility methods", author="simon987", author_email="me@simon987.net", diff --git a/test/test_retry.py b/test/test_retry.py new file mode 100644 index 0000000..e4896b8 --- /dev/null +++ b/test/test_retry.py @@ -0,0 +1,27 @@ +from unittest import TestCase +from hexlib.db import VolatileState, VolatileBooleanState +from hexlib.misc import retry + + +class TestRetry(TestCase): + + def test_simple(self): + @retry(attempts=3) + def a(i): + return i + 1 + + self.assertEqual(a(1), 2) + + def test_error(self): + arr = [] + + def cb(e): + arr.append(e) + + @retry(attempts=3, callback=cb) + def a(i): + raise Exception("err") + + a(1) + + self.assertEqual(3, len(arr))