Merge pull request #35 from k2fong/class-method-support

Class method support

2 files changed, 14 insertions, 1 deletions

diff --git a/tests/test_timeout_decorator.py b/tests/test_timeout_decorator.py
index eac0cb3..d1a0efd 100644
--- a/tests/test_timeout_decorator.py
+++ b/tests/test_timeout_decorator.py
@@ -20,6 +20,15 @@ def test_timeout_decorator_arg(use_signals):
         f()
 
 
+def test_timeout_class_method(use_signals):
+    class c():
+        @timeout(1, use_signals=use_signals)
+        def f(self):
+            time.sleep(2)
+    with pytest.raises(TimeoutError):
+        c().f()
+
+
 def test_timeout_kwargs(use_signals):
     @timeout(3, use_signals=use_signals)
     def f():
diff --git a/timeout_decorator/timeout_decorator.py b/timeout_decorator/timeout_decorator.py
index 466f33d..58802a5 100644
--- a/timeout_decorator/timeout_decorator.py
+++ b/timeout_decorator/timeout_decorator.py
@@ -73,7 +73,11 @@ def timeout(seconds=None, use_signals=True, timeout_exception=TimeoutError):
                         signal.signal(signal.SIGALRM, old)
             return new_function
         else:
-            return _Timeout(function, timeout_exception, seconds)
+            @wraps(function)
+            def new_function(*args, **kwargs):
+                timeout_wrapper = _Timeout(function, timeout_exception, seconds)
+                return timeout_wrapper(*args, **kwargs)
+            return new_function
 
     return decorate