Merge pull request #31 from levi-rs/use-alternative-exception

Add ability to specify what exception gets raised

5 files changed, 33 insertions, 7 deletions

diff --git a/README.rst b/README.rst
index 5573180..e04b76c 100644
--- a/README.rst
+++ b/README.rst
@@ -36,6 +36,23 @@ Usage
     if __name__ == '__main__':
         mytest()
 
+Specify an alternate exception to raise on timeout:
+
+::
+
+    import time
+    import timeout_decorator
+
+    @timeout_decorator.timeout(5, timeout_exception=StopIteration)
+    def mytest():
+        print "Start"
+        for i in range(1,10):
+            time.sleep(1)
+            print "%d seconds have passed" % i
+
+    if __name__ == '__main__':
+        mytest()
+
 Multithreading
 --------------
 
diff --git a/setup.py b/setup.py
index bff1cad..4544700 100644
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,7 @@ long_description = (
 
 setup(
     name='timeout-decorator',
-    version='0.3.2',
+    version='0.3.3',
     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 3196fdb..eac0cb3 100644
--- a/tests/test_timeout_decorator.py
+++ b/tests/test_timeout_decorator.py
@@ -28,6 +28,14 @@ def test_timeout_kwargs(use_signals):
         f(timeout=1)
 
 
+def test_timeout_alternate_exception(use_signals):
+    @timeout(3, use_signals=use_signals, timeout_exception=StopIteration)
+    def f():
+        time.sleep(2)
+    with pytest.raises(StopIteration):
+        f(timeout=1)
+
+
 def test_timeout_no_seconds(use_signals):
     @timeout(use_signals=use_signals)
     def f():
diff --git a/timeout_decorator/__init__.py b/timeout_decorator/__init__.py
index 7a0181b..41c9af9 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.1.0'
+__version__ = '0.3.3'
diff --git a/timeout_decorator/timeout_decorator.py b/timeout_decorator/timeout_decorator.py
index 0f12500..466f33d 100644
--- a/timeout_decorator/timeout_decorator.py
+++ b/timeout_decorator/timeout_decorator.py
@@ -35,7 +35,7 @@ class TimeoutError(AssertionError):
         return repr(self.value)
 
 
-def timeout(seconds=None, use_signals=True):
+def timeout(seconds=None, use_signals=True, timeout_exception=TimeoutError):
     """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 +57,7 @@ def timeout(seconds=None, use_signals=True):
 
         if use_signals:
             def handler(signum, frame):
-                raise TimeoutError()
+                raise timeout_exception()
 
             @wraps(function)
             def new_function(*args, **kwargs):
@@ -73,7 +73,7 @@ def timeout(seconds=None, use_signals=True):
                         signal.signal(signal.SIGALRM, old)
             return new_function
         else:
-            return _Timeout(function, seconds)
+            return _Timeout(function, timeout_exception, seconds)
 
     return decorate
 
@@ -101,10 +101,11 @@ class _Timeout(object):
     to be made and termination of execution after a timeout has passed.
     """
 
-    def __init__(self, function, limit):
+    def __init__(self, function, timeout_exception, limit):
         """Initialize instance in preparation for being called."""
         self.__limit = limit
         self.__function = function
+        self.__timeout_exception = timeout_exception
         self.__name__ = function.__name__
         self.__doc__ = function.__doc__
         self.__timeout = time.time()
@@ -135,7 +136,7 @@ class _Timeout(object):
         """Terminate any possible execution of the embedded function."""
         if self.__process.is_alive():
             self.__process.terminate()
-        raise TimeoutError()
+        raise self.__timeout_exception()
 
     @property
     def ready(self):