From 78a80a16d13a4b7eeb7a4b92f219f95eecf7b361 Mon Sep 17 00:00:00 2001 From: Josh Date: Sat, 10 Jun 2017 03:08:06 -0400 Subject: Can set exception message in decorator --- .vscode/settings.json | 3 +++ tests/test_timeout_decorator.py | 15 +++++++++++++++ timeout_decorator/timeout_decorator.py | 26 ++++++++++++++++++++------ 3 files changed, 38 insertions(+), 6 deletions(-) create mode 100644 .vscode/settings.json diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..fe71598 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.linting.pylintEnabled": false +} \ No newline at end of file diff --git a/tests/test_timeout_decorator.py b/tests/test_timeout_decorator.py index d1a0efd..be6e314 100644 --- a/tests/test_timeout_decorator.py +++ b/tests/test_timeout_decorator.py @@ -86,3 +86,18 @@ def test_timeout_pickle_error(): return Test() with pytest.raises(TimeoutError): f() + +def test_timeout_custom_exception_message(): + message = "Custom fail message" + @timeout(seconds=1, exception_message=message) + def f(): + time.sleep(2) + with pytest.raises(TimeoutError, match=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/timeout_decorator.py b/timeout_decorator/timeout_decorator.py index 58802a5..425f893 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 == 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): @@ -156,4 +170,4 @@ class _Timeout(object): flag, load = self.__queue.get() if flag: return load - raise load + raise load \ No newline at end of file -- cgit v1.2.3