diff --git a/pyhelper_utils/general.py b/pyhelper_utils/general.py index 07d3885..80bf98a 100644 --- a/pyhelper_utils/general.py +++ b/pyhelper_utils/general.py @@ -1,8 +1,9 @@ +from __future__ import annotations import re from time import sleep from functools import wraps from logging import Logger -from typing import Any, Optional +from typing import Any def tts(ts: Any) -> int: @@ -39,14 +40,21 @@ def tts(ts: Any) -> int: return int(ts) -def ignore_exceptions(logger: Optional[Logger] = None, retry: int = 0, retry_interval: int = 1) -> Any: +def ignore_exceptions( + retry: int = 0, + retry_interval: int = 1, + return_on_error: Any = None, + logger: Logger | None = None, +) -> Any: """ Decorator to ignore exceptions with support for retry. Args: - logger (Logger): logger to use, if not passed no logs will be displayed. retry (int): Number of retry if the underline function throw exception. retry_interval (int): Number of seconds to wait between retries. + return_on_error (Any): Return value if the underline function throw exception. + logger (Logger): logger to use, if not passed no logs will be displayed. + Returns: any: the underline function return value. @@ -68,7 +76,7 @@ def inner(*args, **kwargs): if logger: logger.info(f"{func.__name__} error: {ex}") - return None + return return_on_error return inner diff --git a/tests/test_general.py b/tests/test_general.py index 9fd8463..c7e03ec 100644 --- a/tests/test_general.py +++ b/tests/test_general.py @@ -5,7 +5,16 @@ @pytest.fixture def func_for_ignore_exception(): - @ignore_exceptions(logger=logging.getLogger(), retry=1) + @ignore_exceptions(logger=logging.getLogger(), retry=1, retry_interval=1) + def _foo(): + raise ValueError() + + return _foo + + +@pytest.fixture +def func_for_ignore_exception_with_return_value_on_error(): + @ignore_exceptions(logger=logging.getLogger(), retry_interval=1, return_on_error="test") def _foo(): raise ValueError() @@ -20,3 +29,8 @@ def test_tts(): def test_ignore_exceptions(func_for_ignore_exception): assert not isinstance(func_for_ignore_exception(), Exception) + assert not func_for_ignore_exception() + + +def test_ignore_exceptions_with_return_value(func_for_ignore_exception_with_return_value_on_error): + assert func_for_ignore_exception_with_return_value_on_error() == "test"