Merge pull request #36 from joshddunn/master

Set custom exception message in decorator

5 files changed, 38 insertions, 7 deletions

diff --git a/.gitignore b/.gitignore
index c0ffae1..0b7661f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -40,3 +40,4 @@ nosetests.xml
 .pydevproject
 
 .ropeproject/
+.vscode/
diff --git a/setup.py b/setup.py
index 4544700..e001557 100644
--- a/setup.py
+++ b/setup.py
@@ -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):