convert wrap_greenthread from a function to a method of EventletEventLoop

6 files changed, 122 insertions, 117 deletions

diff --git a/aiogreen.py b/aiogreen.py
index 8912d43..a754507 100644
--- a/aiogreen.py
+++ b/aiogreen.py
@@ -282,58 +282,55 @@ class EventletEventLoop(asyncio.SelectorEventLoop):
         future.add_done_callback(done)
         return event.wait()
 
+    def wrap_greenthread(self, gt):
+        """Wrap an eventlet GreenThread, or a greenlet, into a Future object.
 
-class EventletEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
-    _loop_factory = EventletEventLoop
-
-
-def wrap_greenthread(gt, loop=None):
-    """Wrap an eventlet GreenThread, or a greenlet, into a Future object.
-
-    The Future object waits for the completion of a greenthread. The result
-    or the exception of the greenthread will be stored in the Future object.
+        The Future object waits for the completion of a greenthread. The result
+        or the exception of the greenthread will be stored in the Future object.
 
-    The greenthread must be wrapped before its execution starts. If the
-    greenthread is running or already finished, an exception is raised.
+        The greenthread must be wrapped before its execution starts. If the
+        greenthread is running or already finished, an exception is raised.
 
-    For greenlets, the run attribute must be set.
-    """
-    if loop is None:
-        loop = asyncio.get_event_loop()
-    fut = asyncio.Future(loop=loop)
-
-    if not isinstance(gt, greenlet.greenlet):
-        raise TypeError("greenthread or greenlet request, not %s"
-                        % type(gt))
+        For greenlets, the run attribute must be set.
+        """
+        fut = asyncio.Future(loop=self)
+
+        if not isinstance(gt, greenlet.greenlet):
+            raise TypeError("greenthread or greenlet request, not %s"
+                            % type(gt))
+
+        if gt:
+            raise RuntimeError("wrap_greenthread: the greenthread is running")
+        if gt.dead:
+            raise RuntimeError("wrap_greenthread: the greenthread already finished")
+
+        if isinstance(gt, eventlet.greenthread.GreenThread):
+            orig_main = gt.run
+            def wrap_func(*args, **kw):
+                try:
+                    orig_main(*args, **kw)
+                except Exception as exc:
+                    fut.set_exception(exc)
+                else:
+                    result = gt.wait()
+                    fut.set_result(result)
+            gt.run = wrap_func
+        else:
+            try:
+                orig_func = gt.run
+            except AttributeError:
+                raise RuntimeError("wrap_greenthread: the run attribute "
+                                   "of the greenlet is not set")
+            def wrap_func(*args, **kw):
+                try:
+                    result = orig_func(*args, **kw)
+                except Exception as exc:
+                    fut.set_exception(exc)
+                else:
+                    fut.set_result(result)
+            gt.run = wrap_func
+        return fut
 
-    if gt:
-        raise RuntimeError("wrap_greenthread: the greenthread is running")
-    if gt.dead:
-        raise RuntimeError("wrap_greenthread: the greenthread already finished")
 
-    if isinstance(gt, eventlet.greenthread.GreenThread):
-        orig_main = gt.run
-        def wrap_func(*args, **kw):
-            try:
-                orig_main(*args, **kw)
-            except Exception as exc:
-                fut.set_exception(exc)
-            else:
-                result = gt.wait()
-                fut.set_result(result)
-        gt.run = wrap_func
-    else:
-        try:
-            orig_func = gt.run
-        except AttributeError:
-            raise RuntimeError("wrap_greenthread: the run attribute "
-                               "of the greenlet is not set")
-        def wrap_func(*args, **kw):
-            try:
-                result = orig_func(*args, **kw)
-            except Exception as exc:
-                fut.set_exception(exc)
-            else:
-                fut.set_result(result)
-        gt.run = wrap_func
-    return fut
+class EventletEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
+    _loop_factory = EventletEventLoop
diff --git a/doc/changelog.rst b/doc/changelog.rst
index 14a4e6f..be890ca 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -7,6 +7,10 @@ Version 0.4
 * Rename the :class:`EventLoop` class to :class:`EventletEventLoop`,
   and rename the :class:`EventLoopPolicy` class to
   :class:`EventletEventLoopPolicy`
+* Convert :func:`link_future` function to a method of the
+  :class:`EventletEventLoop` class: :meth:`EventletEventLoop.link_future`
+* Convert :func:`wrap_greenthread` function to a method of the
+  :class:`EventletEventLoop` class: :meth:`EventletEventLoop.wrap_greenthread`
 
 2014-10-23: version 0.3
 -----------------------
diff --git a/doc/index.rst b/doc/index.rst
index d61e650..8b4da1f 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -12,8 +12,8 @@ possible to write asyncio code in a project currently written for eventlet.
 
 aiogreen allows to use greenthreads in asyncio coroutines, and to use asyncio
 coroutines, tasks and futures in greenthreads: see
-:meth:`EventletEventLoop.link_future` method and :func:`wrap_greenthread`
-function.
+:meth:`EventletEventLoop.link_future` and
+:meth:`EventletEventLoop.wrap_greenthread` methods.
 
 The main visible difference between aiogreen and trollius is the behaviour of
 ``run_forever()``: ``run_forever()`` blocks with trollius, whereas it runs in a
diff --git a/doc/using.rst b/doc/using.rst
index 827d2ad..cf0a4c1 100644
--- a/doc/using.rst
+++ b/doc/using.rst
@@ -199,64 +199,66 @@ EventletEventLoop
            computation in progress...
            1 + 2 = 3
 
-wrap_greenthread
-----------------
+   .. method:: wrap_greenthread(gt)
 
-.. function:: wrap_greenthread(gt)
+      Wrap an eventlet GreenThread, or a greenlet, 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 result or
+      the exception of the greenthread will be stored in the Future object.
 
-   The Future object waits for the completion of a greenthread. The result or
-   the exception of the greenthread will be stored in the Future object.
+      The greenthread must be wrapped before its execution starts.  If the
+      greenthread is running or already finished, an exception is raised.
 
-   The greenthread must be wrapped before its execution starts.  If the
-   greenthread is running or already finished, an exception is raised.
+      For greenlets, the ``run`` attribute must be set.
 
-   For greenlets, the ``run`` attribute must be set.
+      .. versionchanged:: 0.4
 
-   .. versionchanged:: 0.3
+         The function was converted to a method of the
+         :class:`EventletEventLoop` class.
 
-     An exception is now raised if the greenthread is running or already
-     finished. In debug mode, the exception is not more logged to sys.stderr
-     for greenthreads.
+      .. versionchanged:: 0.3
 
-   Example of trollius coroutine waiting for a greenthread. The ``progress()``
-   callback is called regulary to see that the event loop in not blocked::
+        An exception is now raised if the greenthread is running or already
+        finished. In debug mode, the exception is not more logged to sys.stderr
+        for greenthreads.
 
-        import aiogreen
-        import eventlet
-        import trollius as asyncio
-        from trollius import From, Return
+      Example of trollius coroutine waiting for a greenthread. The ``progress()``
+      callback is called regulary to see that the event loop in not blocked::
 
-        def progress():
-            print("computation in progress...")
-            loop.call_later(0.5, progress)
+           import aiogreen
+           import eventlet
+           import trollius as asyncio
+           from trollius import From, Return
 
-        def slow_sum(x, y):
-            eventlet.sleep(1.0)
-            return x + y
+           def progress():
+               print("computation in progress...")
+               loop.call_later(0.5, progress)
 
-        @asyncio.coroutine
-        def coro_sum():
-            loop.call_soon(progress)
+           def slow_sum(x, y):
+               eventlet.sleep(1.0)
+               return x + y
 
-            gt = eventlet.spawn(slow_sum, 1, 2)
-            fut = aiogreen.wrap_greenthread(gt, loop=loop)
+           @asyncio.coroutine
+           def coro_sum():
+               loop.call_soon(progress)
 
-            result = yield From(fut)
-            print("1 + 2 = %s" % result)
+               gt = eventlet.spawn(slow_sum, 1, 2)
+               fut = aiogreen.wrap_greenthread(gt, loop=loop)
 
-        asyncio.set_event_loop_policy(aiogreen.EventletEventLoopPolicy())
-        loop = asyncio.get_event_loop()
-        loop.run_until_complete(coro_sum())
-        loop.close()
+               result = yield From(fut)
+               print("1 + 2 = %s" % result)
 
-   Output::
+           asyncio.set_event_loop_policy(aiogreen.EventletEventLoopPolicy())
+           loop = asyncio.get_event_loop()
+           loop.run_until_complete(coro_sum())
+           loop.close()
 
-        computation in progress...
-        computation in progress...
-        computation in progress...
-        1 + 2 = 3
+      Output::
+
+           computation in progress...
+           computation in progress...
+           computation in progress...
+           1 + 2 = 3
 
 
 EventletEventLoopPolicy
diff --git a/tests/test_eventlet.py b/tests/test_eventlet.py
index c7ee878..ad5fa6c 100644
--- a/tests/test_eventlet.py
+++ b/tests/test_eventlet.py
@@ -20,20 +20,20 @@ try:
 
     exec('''if 1:
         @asyncio.coroutine
-        def coro_wrap_greenthread():
+        def coro_wrap_greenthread(loop):
             result = []
 
             gt = eventlet.spawn(eventlet_slow_append, result, 1, 0.020)
-            value = yield from aiogreen.wrap_greenthread(gt)
+            value = yield from loop.wrap_greenthread(gt)
             result.append(value)
 
             gt = eventlet.spawn(eventlet_slow_append, result, 2, 0.010)
-            value = yield from aiogreen.wrap_greenthread(gt)
+            value = yield from loop.wrap_greenthread(gt)
             result.append(value)
 
             gt = eventlet.spawn(eventlet_slow_error)
             try:
-                yield from aiogreen.wrap_greenthread(gt)
+                yield from loop.wrap_greenthread(gt)
             except ValueError as exc:
                 result.append(str(exc))
 
@@ -56,20 +56,20 @@ except ImportError:
     from trollius import From, Return
 
     @asyncio.coroutine
-    def coro_wrap_greenthread():
+    def coro_wrap_greenthread(loop):
         result = []
 
         gt = eventlet.spawn(eventlet_slow_append, result, 1, 0.020)
-        value = yield From(aiogreen.wrap_greenthread(gt))
+        value = yield From(loop.wrap_greenthread(gt))
         result.append(value)
 
         gt = eventlet.spawn(eventlet_slow_append, result, 2, 0.010)
-        value = yield From(aiogreen.wrap_greenthread(gt))
+        value = yield From(loop.wrap_greenthread(gt))
         result.append(value)
 
         gt = eventlet.spawn(eventlet_slow_error)
         try:
-            yield From(aiogreen.wrap_greenthread(gt))
+            yield From(loop.wrap_greenthread(gt))
         except ValueError as exc:
             result.append(str(exc))
 
@@ -265,7 +265,7 @@ class WrapGreenthreadTests(tests.TestCase):
             return 'ok'
 
         gt = eventlet.spawn(func)
-        fut = aiogreen.wrap_greenthread(gt)
+        fut = self.loop.wrap_greenthread(gt)
         result = self.loop.run_until_complete(fut)
         self.assertEqual(result, 'ok')
 
@@ -278,7 +278,7 @@ class WrapGreenthreadTests(tests.TestCase):
         # FIXME: the unit test must fail!?
         with tests.mock.patch('traceback.print_exception') as print_exception:
             gt = eventlet.spawn(func)
-            fut = aiogreen.wrap_greenthread(gt)
+            fut = self.loop.wrap_greenthread(gt)
             self.assertRaises(ValueError, self.loop.run_until_complete, fut)
 
         # the exception must not be logger by traceback: the caller must
@@ -287,7 +287,7 @@ class WrapGreenthreadTests(tests.TestCase):
 
     def test_wrap_greenthread_running(self):
         def func():
-            return aiogreen.wrap_greenthread(gt)
+            return self.loop.wrap_greenthread(gt)
 
         self.loop.set_debug(False)
         gt = eventlet.spawn(func)
@@ -304,23 +304,23 @@ class WrapGreenthreadTests(tests.TestCase):
 
         msg = "wrap_greenthread: the greenthread already finished"
         self.assertRaisesRegexp(RuntimeError, msg,
-                                aiogreen.wrap_greenthread, gt)
+                                self.loop.wrap_greenthread, gt)
 
     def test_coro_wrap_greenthread(self):
-        result = self.loop.run_until_complete(coro_wrap_greenthread())
+        result = self.loop.run_until_complete(coro_wrap_greenthread(self.loop))
         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)
+        self.assertRaises(TypeError, self.loop.wrap_greenthread, func)
 
         @asyncio.coroutine
         def coro_func():
             pass
         coro_obj = coro_func()
         self.addCleanup(coro_obj.close)
-        self.assertRaises(TypeError, aiogreen.wrap_greenthread, coro_obj)
+        self.assertRaises(TypeError, self.loop.wrap_greenthread, coro_obj)
 
 
 class WrapGreenletTests(tests.TestCase):
@@ -330,7 +330,7 @@ class WrapGreenletTests(tests.TestCase):
             return "ok"
 
         gt = eventlet.spawn_n(func)
-        fut = aiogreen.wrap_greenthread(gt)
+        fut = self.loop.wrap_greenthread(gt)
         result = self.loop.run_until_complete(fut)
         self.assertEqual(result, "ok")
 
@@ -341,7 +341,7 @@ class WrapGreenletTests(tests.TestCase):
             raise ValueError(7)
 
         gt = eventlet.spawn_n(func)
-        fut = aiogreen.wrap_greenthread(gt)
+        fut = self.loop.wrap_greenthread(gt)
         self.assertRaises(ValueError, self.loop.run_until_complete, fut)
 
     def test_wrap_greenlet_running(self):
@@ -350,7 +350,7 @@ class WrapGreenletTests(tests.TestCase):
         def func():
             try:
                 gt = eventlet.getcurrent()
-                fut = aiogreen.wrap_greenthread(gt)
+                fut = self.loop.wrap_greenthread(gt)
             except Exception as exc:
                 event.send_exception(exc)
             else:
@@ -368,7 +368,8 @@ class WrapGreenletTests(tests.TestCase):
         gt = eventlet.spawn_n(func)
         event.wait()
         msg = "wrap_greenthread: the greenthread already finished"
-        self.assertRaisesRegexp(RuntimeError, msg, aiogreen.wrap_greenthread, gt)
+        self.assertRaisesRegexp(RuntimeError, msg,
+                                self.loop.wrap_greenthread, gt)
 
 
 if __name__ == '__main__':
diff --git a/tests/test_greenlet.py b/tests/test_greenlet.py
index f411a8e..c3f1073 100644
--- a/tests/test_greenlet.py
+++ b/tests/test_greenlet.py
@@ -9,7 +9,7 @@ class WrapGreenletTests(tests.TestCase):
             return value * 3
 
         gl = greenlet.greenlet(func)
-        fut = aiogreen.wrap_greenthread(gl)
+        fut = self.loop.wrap_greenthread(gl)
         gl.switch(5)
         result = self.loop.run_until_complete(fut)
         self.assertEqual(result, 15)
@@ -19,7 +19,7 @@ class WrapGreenletTests(tests.TestCase):
             raise ValueError(7)
 
         gl = greenlet.greenlet(func)
-        fut = aiogreen.wrap_greenthread(gl)
+        fut = self.loop.wrap_greenthread(gl)
         gl.switch()
         self.assertRaises(ValueError, self.loop.run_until_complete, fut)
 
@@ -27,12 +27,12 @@ class WrapGreenletTests(tests.TestCase):
         gl = greenlet.greenlet()
         msg = "wrap_greenthread: the run attribute of the greenlet is not set"
         self.assertRaisesRegexp(RuntimeError, msg,
-                                aiogreen.wrap_greenthread, gl)
+                                self.loop.wrap_greenthread, gl)
 
     def test_wrap_greenlet_running(self):
         def func(value):
             gl = greenlet.getcurrent()
-            return aiogreen.wrap_greenthread(gl)
+            return self.loop.wrap_greenthread(gl)
 
         gl = greenlet.greenlet(func)
         msg = "wrap_greenthread: the greenthread is running"
@@ -45,7 +45,8 @@ class WrapGreenletTests(tests.TestCase):
         gl = greenlet.greenlet(func)
         gl.switch(5)
         msg = "wrap_greenthread: the greenthread already finished"
-        self.assertRaisesRegexp(RuntimeError, msg, aiogreen.wrap_greenthread, gl)
+        self.assertRaisesRegexp(RuntimeError, msg,
+                                self.loop.wrap_greenthread, gl)
 
 
 if __name__ == '__main__':