Merge pull request #1 from pfultz2/gyield-timeout

Add optional timeout parameter to gyield

greenado/__init__.py

@@ -1,3 +1,3 @@
 
-from .concurrent import gcall, generator, groutine, gyield
+from .concurrent import gcall, generator, groutine, gyield, TimeoutError
 from .version import __version__

greenado/concurrent.py

@@ -23,6 +23,8 @@
 from tornado import concurrent, gen
 from tornado.ioloop import IOLoop
 
+class TimeoutError(Exception):
+    """Exception raised by ``gyield`` in timeout."""
 
 def gcall(f, *args, **kwargs):
     '''
@@ -170,7 +172,7 @@ def greenlet_base():
     return wrapper
 
 
-def gyield(future):
+def gyield(future, timeout=None):
     '''
         This is functionally equivalent to the 'yield' statements used in a
         :func:`@gen.coroutine <tornado.gen.coroutine>`, but doesn't require
@@ -197,8 +199,14 @@ def gyield(future):
 
     # don't switch/wait if the future is already ready to go
     if not future.done():
-
+        timeout_handle = None
+        if timeout != None and timeout > 0:
+            timeout_handle = IOLoop.current().add_timeout(
+                IOLoop.current().time() + timeout,
+                lambda: future.set_exception(TimeoutError("Timeout")))
+
         def on_complete(result):
+            if timeout_handle != None: IOLoop.current().remove_timeout(timeout_handle)
             gr.switch()
 
         IOLoop.current().add_future(future, on_complete)

tests/test_greenado.py

@@ -5,6 +5,7 @@
 
 from tornado import gen, concurrent
 from tornado.ioloop import IOLoop
+import time
 
 
 class DummyException(Exception):
@@ -146,6 +147,27 @@ def _main():
 
     main_result = IOLoop.current().run_sync(_main)
     assert main_result == True
+
+def test_gyield_timeout():
+    '''Ensure that timout throws an exception after time runs out'''
+
+    future = concurrent.Future()
+
+    def _inner():
+        try:
+            greenado.gyield(future, 2)
+            return 0
+        except greenado.TimeoutError:
+            return 1
+
+    @greenado.groutine
+    def _main():
+        return _inner() + 1
+
+    start_time = time.time()
+    main_retval = IOLoop.current().run_sync(_main)
+    assert main_retval == 2
+    assert time.time() > start_time + 2
 
 
 def test_generator_error():