wrap_greenthread() now also works on greenlet objects

5 files changed, 128 insertions, 15 deletions

diff --git a/aiogreen.py b/aiogreen.py
index b5d9647..8fc643e 100644
--- a/aiogreen.py
+++ b/aiogreen.py
@@ -1,4 +1,5 @@
 import eventlet.hubs.hub
+import greenlet
 import sys
 socket = eventlet.patcher.original('socket')
 threading = eventlet.patcher.original('threading')
@@ -234,20 +235,41 @@ class EventLoopPolicy(asyncio.DefaultEventLoopPolicy):
 
 
 def wrap_greenthread(gt, loop=None):
-    """Wrap an eventlet greenthread into a Future object."""
+    """Wrap an eventlet GreenThread or a greenlet into a Future object.
+
+    The greenlet must not be running."""
     if loop is None:
         loop = asyncio.get_event_loop()
     fut = asyncio.Future(loop=loop)
 
-    def copy_result(gt):
-        try:
-            value = gt.wait()
-        except Exception as exc:
-            loop.call_soon(fut.set_exception, exc)
-        else:
-            loop.call_soon(fut.set_result, value)
-
-    gt.link(copy_result)
+    if isinstance(gt, eventlet.greenthread.GreenThread):
+        def copy_result(gt):
+            try:
+                result = gt.wait()
+            except Exception as exc:
+                loop.call_soon(fut.set_exception, exc)
+            else:
+                loop.call_soon(fut.set_result, result)
+
+        gt.link(copy_result)
+    elif isinstance(gt, greenlet.greenlet):
+        if gt:
+            raise RuntimeError("cannot wrap a running greenlet")
+        if gt.dead:
+            raise RuntimeError("cannot wrap a greenlet which already finished")
+
+        orig_func = gt.run
+        def wrap_func(*args, **kw):
+            try:
+                result = orig_func(*args, **kw)
+            except Exception as exc:
+                loop.call_soon(fut.set_exception, exc)
+            else:
+                loop.call_soon(fut.set_result, result)
+        gt.run = wrap_func
+    else:
+        raise TypeError("greenthread or greenlet request, not %s"
+                        % type(gt))
     return fut
 
 
diff --git a/doc/changelog.rst b/doc/changelog.rst
index b7893e9..5853cce 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -1,6 +1,11 @@
 Changelog
 =========
 
+Version 0.3 (development version)
+---------------------------------
+
+* :func:`wrap_greenthread` now also works on greenlet objects.
+
 2014-10-21: version 0.2
 -----------------------
 
diff --git a/doc/status.rst b/doc/status.rst
index 7b0802d..21e4f71 100644
--- a/doc/status.rst
+++ b/doc/status.rst
@@ -1,6 +1,8 @@
 To do
 =====
 
+* doesn't make sense to call wrap_greenthread() from the running greenthread?
+  (it works for greenthreads, but it doesn't work for greenlets)
 * wrap_greenthread() must not log the exception to sys.stderr if the
   greenthread didn't start
 * register signals in eventlet hub, only needed for pyevent hub?
diff --git a/doc/using.rst b/doc/using.rst
index 25687c6..8e62110 100644
--- a/doc/using.rst
+++ b/doc/using.rst
@@ -134,10 +134,12 @@ aiogreen specific functions:
 
 .. function:: wrap_greenthread(gt)
 
-   Wrap a greenthread into a Future object.
+   Wrap an eventlet GreenThread or a greenlet into a Future object.
 
    The Future object waits for the completion of a greenthread.
 
+   The greenlet must not be running.
+
    In debug mode, if the greenthread raises an exception, the exception is
    logged to ``sys.stderr`` by eventlet, even if the exception is copied to the
    Future object.
diff --git a/tests/test_eventlet.py b/tests/test_eventlet.py
index 14e4483..8105cf2 100644
--- a/tests/test_eventlet.py
+++ b/tests/test_eventlet.py
@@ -162,10 +162,6 @@ class EventletTests(tests.TestCase):
         self.loop.run_forever()
         self.assertEqual(result, ["spawn", "spawn_after"])
 
-    def test_coro_wrap_greenthread(self):
-        result = self.loop.run_until_complete(coro_wrap_greenthread())
-        self.assertEqual(result, [1, 10, 2, 20, 'error', 4])
-
     def test_greenthread_link_future(self):
         result = []
         self.loop.call_soon(eventlet.spawn,
@@ -174,6 +170,92 @@ class EventletTests(tests.TestCase):
         self.assertEqual(result, [1, 10, 2, 20, 'error', 4])
 
 
+class WrapGreenthreadTests(tests.TestCase):
+    def test_wrap_greenthread(self):
+        def func():
+            eventlet.sleep(0.010)
+            return 'ok'
+
+        gt = eventlet.spawn(func)
+        fut = aiogreen.wrap_greenthread(gt)
+        result = self.loop.run_until_complete(fut)
+        self.assertEqual(result, 'ok')
+
+    def test_wrap_greenthread_running(self):
+        event = eventlet.event.Event()
+
+        def func():
+            return aiogreen.wrap_greenthread(gt)
+
+        gt = eventlet.spawn(func)
+        fut1 = aiogreen.wrap_greenthread(gt)
+        fut2 = self.loop.run_until_complete(fut1)
+        fut3 = self.loop.run_until_complete(fut2)
+        self.assertIs(fut3, fut2)
+
+    def test_wrap_greenthread_dead(self):
+        def func():
+            return 'ok'
+
+        gt = eventlet.spawn(func)
+        result = gt.wait()
+        self.assertEqual(result, 'ok')
+
+        fut = aiogreen.wrap_greenthread(gt)
+        result = self.loop.run_until_complete(fut)
+        self.assertEqual(result, 'ok')
+
+    def test_coro_wrap_greenthread(self):
+        result = self.loop.run_until_complete(coro_wrap_greenthread())
+        self.assertEqual(result, [1, 10, 2, 20, 'error', 4])
+
+    def test_wrap_invalid_type(self):
+        def func():
+            pass
+        self.assertRaises(TypeError, aiogreen.wrap_greenthread, func)
+
+        @asyncio.coroutine
+        def coro_func():
+            pass
+        coro_obj = coro_func()
+        self.assertRaises(TypeError, aiogreen.wrap_greenthread, coro_obj)
+
+
+class WrapGreenletTests(tests.TestCase):
+    def test_wrap_greenlet(self):
+        def func():
+            eventlet.sleep(0.010)
+            return "ok"
+        gt = eventlet.spawn_n(func)
+        fut = aiogreen.wrap_greenthread(gt)
+        result = self.loop.run_until_complete(fut)
+        self.assertEqual(result, "ok")
+
+    def test_wrap_greenlet_running(self):
+        event = eventlet.event.Event()
+
+        def func():
+            try:
+                gt = eventlet.getcurrent()
+                fut = aiogreen.wrap_greenthread(gt)
+            except Exception as exc:
+                event.send_exception(exc)
+            else:
+                event.send(fut)
+
+        gt = eventlet.spawn_n(func)
+        self.assertRaises(RuntimeError, event.wait)
+
+    def test_wrap_greenlet_dead(self):
+        event = eventlet.event.Event()
+        def func():
+            event.send('done')
+        gt = eventlet.spawn_n(func)
+        event.wait()
+
+        self.assertRaises(RuntimeError, aiogreen.wrap_greenthread, gt)
+
+
 if __name__ == '__main__':
     import unittest
     unittest.main()