diff options
-rw-r--r-- | .gitignore | 1 | ||||
-rw-r--r-- | setup.py | 2 | ||||
-rw-r--r-- | tests/test_timeout_decorator.py | 16 | ||||
-rw-r--r-- | timeout_decorator/__init__.py | 2 | ||||
-rw-r--r-- | timeout_decorator/timeout_decorator.py | 24 |
5 files changed, 38 insertions, 7 deletions
@@ -40,3 +40,4 @@ nosetests.xml .pydevproject .ropeproject/ +.vscode/ @@ -27,7 +27,7 @@ long_description = ( setup( name='timeout-decorator', - version='0.3.3', + version='0.4.0', description='Timeout decorator', long_description=long_description, author='Patrick Ng', diff --git a/tests/test_timeout_decorator.py b/tests/test_timeout_decorator.py index d1a0efd..536bc8f 100644 --- a/tests/test_timeout_decorator.py +++ b/tests/test_timeout_decorator.py @@ -86,3 +86,19 @@ def test_timeout_pickle_error(): return Test() with pytest.raises(TimeoutError): f() + + +def test_timeout_custom_exception_message(): + @timeout(seconds=1, exception_message="Custom fail message") + def f(): + time.sleep(2) + with pytest.raises(TimeoutError, match="Custom fail message"): + f() + + +def test_timeout_default_exception_message(): + @timeout(seconds=1) + def f(): + time.sleep(2) + with pytest.raises(TimeoutError, match="Timed Out"): + f() diff --git a/timeout_decorator/__init__.py b/timeout_decorator/__init__.py index 41c9af9..c802c55 100644 --- a/timeout_decorator/__init__.py +++ b/timeout_decorator/__init__.py @@ -4,4 +4,4 @@ from .timeout_decorator import timeout from .timeout_decorator import TimeoutError __title__ = 'timeout_decorator' -__version__ = '0.3.3' +__version__ = '0.4.0' diff --git a/timeout_decorator/timeout_decorator.py b/timeout_decorator/timeout_decorator.py index 58802a5..f05da0f 100644 --- a/timeout_decorator/timeout_decorator.py +++ b/timeout_decorator/timeout_decorator.py @@ -35,7 +35,19 @@ class TimeoutError(AssertionError): return repr(self.value) -def timeout(seconds=None, use_signals=True, timeout_exception=TimeoutError): +def _raise_exception(exception, exception_message): + """ This function checks if a exception message is given. + + If there is no exception message, the default behaviour is maintained. + If there is an exception message, the message is passed to the exception with the 'value' keyword. + """ + if exception_message is None: + raise exception() + else: + raise exception(value=exception_message) + + +def timeout(seconds=None, use_signals=True, timeout_exception=TimeoutError, exception_message=None): """Add a timeout parameter to a function and return it. :param seconds: optional time limit in seconds or fractions of a second. If None is passed, no timeout is applied. @@ -57,7 +69,7 @@ def timeout(seconds=None, use_signals=True, timeout_exception=TimeoutError): if use_signals: def handler(signum, frame): - raise timeout_exception() + _raise_exception(timeout_exception, exception_message) @wraps(function) def new_function(*args, **kwargs): @@ -75,7 +87,7 @@ def timeout(seconds=None, use_signals=True, timeout_exception=TimeoutError): else: @wraps(function) def new_function(*args, **kwargs): - timeout_wrapper = _Timeout(function, timeout_exception, seconds) + timeout_wrapper = _Timeout(function, timeout_exception, exception_message, seconds) return timeout_wrapper(*args, **kwargs) return new_function @@ -105,11 +117,12 @@ class _Timeout(object): to be made and termination of execution after a timeout has passed. """ - def __init__(self, function, timeout_exception, limit): + def __init__(self, function, timeout_exception, exception_message, limit): """Initialize instance in preparation for being called.""" self.__limit = limit self.__function = function self.__timeout_exception = timeout_exception + self.__exception_message = exception_message self.__name__ = function.__name__ self.__doc__ = function.__doc__ self.__timeout = time.time() @@ -140,7 +153,8 @@ class _Timeout(object): """Terminate any possible execution of the embedded function.""" if self.__process.is_alive(): self.__process.terminate() - raise self.__timeout_exception() + + _raise_exception(self.__timeout_exception, self.__exception_message) @property def ready(self): |