feat(util): add retry decorator

podmaker/util/__init__.py

@@ -1,3 +1,4 @@
-__all__ = ['exit_signal', 'ExitSignalError']
+__all__ = ['exit_signal', 'ExitSignalError', 'retry']
 
 from podmaker.util.exit import ExitSignalError, exit_signal
+from podmaker.util.retry import retry

podmaker/util/retry.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import logging
+import sys
+import time
+from datetime import timedelta
+from typing import Callable, Tuple, Type, TypeVar
+
+if sys.version_info < (3, 10):
+    from typing_extensions import ParamSpec
+else:
+    from typing import ParamSpec
+
+
+P = ParamSpec('P')
+T = TypeVar('T')
+_logger = logging.getLogger(__name__)
+
+
+def retry(
+        cnt: int,
+        *,
+        wait: timedelta = timedelta(seconds=0),
+        catch: Type[Exception] | Tuple[Type[Exception], ...] = Exception,
+        logger: logging.Logger = _logger,
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    """
+    A decorator to retry the function when exception raised.
+    The function will be called at least once and at most cnt + 1 times.
+
+    :param cnt: retry count
+    :param wait: wait time between retries
+    :param catch: the exception to retry
+    :param logger: logger to log retry info
+    """
+    if cnt <= 0:
+        raise ValueError('cnt must be positive')
+    wait_seconds = wait.total_seconds()
+
+    def deco(func: Callable[P, T]) -> Callable[P, T]:
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            for _ in range(cnt):
+                try:
+                    return func(*args, **kwargs)
+                except catch:
+                    logger.warning('retrying...')
+                    if wait_seconds > 0:
+                        logger.warning(f'wait {wait_seconds}s before retry')
+                        time.sleep(wait_seconds)
+            return func(*args, **kwargs)
+        return wrapper
+    return deco

tests/provider/test_youtube.py

@@ -9,7 +9,7 @@
 from podmaker.config import OwnerConfig, SourceConfig
 from podmaker.fetcher import YouTube
 from podmaker.storage import ObjectInfo, Storage
-from tests.util import network_available
+from tests.helper import network_available
 
 if sys.version_info >= (3, 11):
     pass

tests/util/test_retry.py

@@ -0,0 +1,35 @@
+import unittest
+from unittest import mock
+
+from podmaker.util import retry
+
+
+class TestRetry(unittest.TestCase):
+    def test_no_exception(self) -> None:
+        spy = mock.Mock(return_value=1)
+        func = retry(3)(spy)
+        self.assertEqual(1, func())
+        self.assertEqual(1, spy.call_count)
+
+    def test_retry_success(self) -> None:
+        spy = mock.Mock(side_effect=[Exception, 1])
+        func = retry(3)(spy)
+        self.assertEqual(1, func())
+        self.assertEqual(2, spy.call_count)
+
+    def test_retry_failed(self) -> None:
+        spy = mock.Mock(side_effect=Exception)
+        func = retry(3)(spy)
+        self.assertRaises(Exception, func)
+        self.assertEqual(4, spy.call_count)
+
+    def test_specify_exception(self) -> None:
+        spy = mock.Mock(side_effect=ValueError)
+        func = retry(3, catch=TypeError)(spy)
+        self.assertRaises(ValueError, func)
+        self.assertEqual(1, spy.call_count)
+
+        spy = mock.Mock(side_effect=ValueError)
+        func = retry(3, catch=ValueError)(spy)
+        self.assertRaises(ValueError, func)
+        self.assertEqual(4, spy.call_count)