Merge pull request #3 from jonblack/feature-timeout-kwarg

Add support for timeout in function kwargs

4 files changed, 50 insertions, 12 deletions

diff --git a/.gitignore b/.gitignore
index bd77452..20680a8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -36,3 +36,5 @@ nosetests.xml
 .mr.developer.cfg
 .project
 .pydevproject
+
+.ropeproject/
diff --git a/requirements b/requirements
new file mode 100644
index 0000000..d095267
--- /dev/null
+++ b/requirements
@@ -0,0 +1,3 @@
+argparse==1.2.1
+nose==1.3.4
+wsgiref==0.1.2
diff --git a/tests/test_timeout_decorator.py b/tests/test_timeout_decorator.py
index 35d0e3e..143dfc8 100644
--- a/tests/test_timeout_decorator.py
+++ b/tests/test_timeout_decorator.py
@@ -2,14 +2,38 @@
 # -*- coding: utf-8 -*-
 
 import time
-import timeout_decorator 
 
-@timeout_decorator.timeout(5)
-def mytest():
-    print "Start"
-    for i in range(1,10):
-        time.sleep(1)
-        print "%d seconds have passed" % i
+from nose.tools import raises
+
+from timeout_decorator import timeout, TimeoutError
+
+
+@raises(TimeoutError)
+def test_timeout_decorator_arg():
+    @timeout(1)
+    def f():
+        time.sleep(2)
+    f()
+
 
-if __name__ == '__main__':
-    mytest()
+@raises(TimeoutError)
+def test_timeout_kwargs():
+    @timeout()
+    def f(timeout):
+        time.sleep(2)
+    f(timeout=1)
+
+
+@raises(ValueError)
+def test_timeout_no_seconds():
+    @timeout()
+    def f(timeout):
+        time.sleep(2)
+    f()
+
+
+def test_timeout_ok():
+    @timeout(seconds=2)
+    def f():
+        time.sleep(1)
+    f()
diff --git a/timeout_decorator/timeout_decorator.py b/timeout_decorator/timeout_decorator.py
index 9876ad6..2b86bfb 100644
--- a/timeout_decorator/timeout_decorator.py
+++ b/timeout_decorator/timeout_decorator.py
@@ -18,19 +18,28 @@ import signal
 
 #http://www.saltycrane.com/blog/2010/04/using-python-timeout-decorator-uploading-s3/
 
+
 class TimeoutError(Exception):
-    def __init__(self, value = "Timed Out"):
+    def __init__(self, value="Timed Out"):
         self.value = value
+
     def __str__(self):
         return repr(self.value)
 
-def timeout(seconds_before_timeout):
+
+def timeout(seconds=None):
     def decorate(f):
         def handler(signum, frame):
             raise TimeoutError()
+
         def new_f(*args, **kwargs):
             old = signal.signal(signal.SIGALRM, handler)
-            signal.alarm(seconds_before_timeout)
+
+            new_seconds = kwargs['timeout'] if 'timeout' in kwargs else seconds
+            if new_seconds is None:
+                raise ValueError("You must provide a timeout value")
+
+            signal.alarm(new_seconds)
             try:
                 result = f(*args, **kwargs)
             finally: