Port greenio to Trollius

greenio/__init__.py

@@ -9,9 +9,46 @@
 
 
 import greenlet
+import sys
+
+try:
+    import trollius
+except ImportError:
+    trollius = None
+try:
+    import asyncio
+except ImportError:
+    if trollius is not None:
+        asyncio = trollius
+    else:
+        raise
+
+if trollius is not None:
+    def _async(future):
+        # trollius iscoroutine() accepts trollius and asyncio coroutine
+        # objects
+        if trollius.iscoroutine(future):
+            loop = asyncio.get_event_loop()
+            return loop.create_task(future)
+        else:
+            return future
+
+    def _create_task(coro):
+        loop = asyncio.get_event_loop()
+        return loop.create_task(coro)
+else:
+    def _async(future):
+        if asyncio.iscoroutine(future):
+            return GreenTask(future)
+        else:
+            return future
 
-import asyncio
-from asyncio import unix_events, tasks, futures
+    def _create_task(coro):
+        return GreenTask(coro)
+
+_FUTURE_CLASSES = (asyncio.Future,)
+if trollius is not None and trollius is not asyncio:
+    _FUTURE_CLASSES += (trollius.Future,)
 
 import sys
 
@@ -28,16 +65,16 @@ class _TaskGreenlet(greenlet.greenlet):
     ``@greenio.task`` is executed in this greenlet"""
 
 
-class GreenTask(asyncio.Task):
+class _GreenTaskMixin(object):
     def __init__(self, *args, **kwargs):
         self._greenlet = None
-        super(GreenTask, self).__init__(*args, **kwargs)
+        super(_GreenTaskMixin, self).__init__(*args, **kwargs)
 
     def _step(self, value=None, exc=None):
         if self._greenlet is None:
             # Means that the task is not currently in a suspended greenlet
             # waiting for results for "yield_from"
-            ovr = super(GreenTask, self)._step
+            ovr = super(_GreenTaskMixin, self)._step
             self._greenlet = _TaskGreenlet(ovr)
 
             # Store a reference to the current task for "yield_from"
@@ -93,21 +130,43 @@ def run_forever(self, *args, **kwargs):
         return self._green_run(ovr, args, kwargs)
 
 
-class GreenUnixSelectorLoop(_GreenLoopMixin, asyncio.SelectorEventLoop):
+class GreenTask(_GreenTaskMixin, asyncio.Task):
     pass
 
 
-class GreenEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
+class GreenUnixSelectorLoop(_GreenLoopMixin, asyncio.SelectorEventLoop):
+    def create_task(self, coro):
+        return GreenTask(coro)
 
+
+class GreenEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
     def new_event_loop(self):
         return GreenUnixSelectorLoop()
 
 
+if trollius is not None:
+    if trollius is not asyncio:
+        class GreenTrolliusTask(_GreenTaskMixin, trollius.Task):
+            pass
+
+        class GreenTrolliusUnixSelectorLoop(_GreenLoopMixin, trollius.SelectorEventLoop):
+            def create_task(self, coro):
+                return GreenTrolliusTask(coro)
+
+        class GreenTrolliusEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
+
+            def new_event_loop(self):
+                return GreenTrolliusUnixSelectorLoop()
+    else:
+        GreenTrolliusTask = GreenTask
+        GreenTrolliusUnixSelectorLoop = GreenUnixSelectorLoop
+        GreenTrolliusEventLoopPolicy = GreenEventLoopPolicy
+
+
 def yield_from(future):
     """A function to use instead of ``yield from`` statement."""
 
-    if asyncio.iscoroutine(future):
-        future = GreenTask(future)
+    future = _async(future)
 
     gl = greenlet.getcurrent()
 
@@ -126,7 +185,7 @@ def yield_from(future):
 
     task = gl.task
 
-    if not isinstance(future, futures.Future):
+    if not isinstance(future, _FUTURE_CLASSES):
         raise RuntimeError(
             'greenlet.yield_from was supposed to receive only Futures, '
             'got {!r} in task {!r}'.format(future, task))
@@ -149,10 +208,11 @@ def task(func):
     """A decorator, allows use of ``yield_from`` in the decorated or
     subsequent coroutines."""
 
-    coro = asyncio.coroutine(func)
+    coro_func = asyncio.coroutine(func)
 
     def task_wrapper(*args, **kwds):
-        return GreenTask(coro(*args, **kwds))
+        coro_obj = coro_func(*args, **kwds)
+        return _create_task(coro_obj)
 
     return task_wrapper
 

greenio/socket.py

@@ -9,7 +9,10 @@
 from them.
 """
 from __future__ import absolute_import
-import asyncio
+try:
+    import asyncio
+except ImportError:
+    import trollius as asyncio
 from socket import error, SOCK_STREAM
 from socket import socket as std_socket
 

tests/test_socket.py

@@ -4,16 +4,21 @@
 ##
 
 
-import asyncio
-import unittest
+try:
+    import asyncio
+    from unittest import TestCase
+except ImportError:
+    import trollius as asyncio
+    from trollius.test_utils import TestCase
+import socket
 
 import greenio
 import greenio.socket as greensocket
 
 import socket as std_socket
 
 
-class SocketTests(unittest.TestCase):
+class SocketTests(TestCase):
 
     def setUp(self):
         asyncio.set_event_loop_policy(greenio.GreenEventLoopPolicy())

tests/test_tasks.py

@@ -5,11 +5,18 @@
 
 
 import greenio
-import asyncio
-import unittest
-
-
-class TaskTests(unittest.TestCase):
+try:
+    trollius = None
+    import asyncio
+    from unittest import TestCase
+except ImportError:
+    import trollius
+    from trollius import From, Return
+    asyncio = trollius
+    from trollius.test_utils import TestCase
+
+
+class TaskTests(TestCase):
     def setUp(self):
         asyncio.set_event_loop_policy(greenio.GreenEventLoopPolicy())
         self.loop = asyncio.new_event_loop()
@@ -19,73 +26,131 @@ def tearDown(self):
         self.loop.close()
         asyncio.set_event_loop_policy(None)
 
-    def test_task_yield_from_plain(self):
-        @asyncio.coroutine
-        def bar():
-            yield
-            return 30
-
-        @asyncio.coroutine
-        def foo():
-            bar_result = greenio.yield_from(asyncio.Task(bar()))
-            return bar_result + 12
-
-        @greenio.task
-        def test():
-            return (yield from foo())
+    if trollius is not None:
+        def test_task_yield_from_plain(self):
+            @asyncio.coroutine
+            def bar():
+                yield From(None)
+                raise Return(30)
+
+            @asyncio.coroutine
+            def foo():
+                bar_result = greenio.yield_from(bar())
+                return bar_result + 12
+
+            @greenio.task
+            def test():
+                res = yield From(foo())
+                raise Return(res)
+
+            fut = test()
+            self.loop.run_until_complete(fut)
+
+            self.assertEqual(fut.result(), 42)
+
+        def test_task_yield_from_exception_propagation(self):
+            non_local = {'CHK': 0}
+
+            @asyncio.coroutine
+            def bar():
+                yield From(None)
+                yield From(None)
+                1/0
+
+            @greenio.task
+            def foo():
+                greenio.yield_from(bar())
+
+            @asyncio.coroutine
+            def test():
+                try:
+                    res = yield From(foo())
+                    raise Return(res)
+                except ZeroDivisionError:
+                    non_local['CHK'] += 1
+
+            self.loop.run_until_complete(test())
+            self.assertEqual(non_local['CHK'], 1)
+
+        def test_task_yield_from_coroutine(self):
+            @asyncio.coroutine
+            def bar():
+                yield From(None)
+                raise Return(5)
+
+            @greenio.task
+            def foo():
+                return greenio.yield_from(bar())
+
+            fut = foo()
+            self.loop.run_until_complete(fut)
+            self.assertEqual(fut.result(), 5)
+    else:
+        exec('''if 1:
+        def test_task_yield_from_plain(self):
+            @asyncio.coroutine
+            def bar():
+                yield from []
+                return 30
+
+            @asyncio.coroutine
+            def foo():
+                bar_result = greenio.yield_from(bar())
+                return bar_result + 12
+
+            @greenio.task
+            def test():
+                return (yield from foo())
 
-        fut = test()
-        self.loop.run_until_complete(fut)
+            fut = test()
+            self.loop.run_until_complete(fut)
 
-        self.assertEqual(fut.result(), 42)
+            self.assertEqual(fut.result(), 42)
 
-    def test_task_yield_from_exception_propagation(self):
-        CHK = 0
+        def test_task_yield_from_exception_propagation(self):
+            non_local = {'CHK': 0}
 
-        @asyncio.coroutine
-        def bar():
-            yield
-            yield
-            1/0
+            @asyncio.coroutine
+            def bar():
+                yield
+                yield
+                1/0
 
-        @greenio.task
-        def foo():
-            greenio.yield_from(asyncio.Task(bar()))
+            @greenio.task
+            def foo():
+                greenio.yield_from(bar())
 
-        @asyncio.coroutine
-        def test():
-            try:
-                return (yield from foo())
-            except ZeroDivisionError:
-                nonlocal CHK
-                CHK += 1
+            @asyncio.coroutine
+            def test():
+                try:
+                    return (yield from foo())
+                except ZeroDivisionError:
+                    non_local['CHK'] += 1
 
-        self.loop.run_until_complete(test())
-        self.assertEqual(CHK, 1)
+            self.loop.run_until_complete(test())
+            self.assertEqual(non_local['CHK'], 1)
 
-    def test_task_yield_from_coroutine(self):
-        @asyncio.coroutine
-        def bar():
-            yield from []
-            return 5
+        def test_task_yield_from_coroutine(self):
+            @asyncio.coroutine
+            def bar():
+                yield from []
+                return 5
 
-        @greenio.task
-        def foo():
-            return greenio.yield_from(bar())
+            @greenio.task
+            def foo():
+                return greenio.yield_from(bar())
 
-        fut = foo()
-        self.loop.run_until_complete(fut)
-        self.assertEqual(fut.result(), 5)
+            fut = foo()
+            self.loop.run_until_complete(fut)
+            self.assertEqual(fut.result(), 5)
+        ''')
 
     def test_task_yield_from_invalid(self):
-        def bar():
-            pass
+        err_msg = 'greenlet.yield_from was supposed to receive only Futures, '
 
         @asyncio.coroutine
         def foo():
-            with self.assertRaisesRegex(
-                    RuntimeError,
-                    '"greenio.yield_from" was supposed to be called'):
-                greenio.yield_from(bar)
+            with self.assertRaisesRegex(RuntimeError, err_msg):
+                greenio.yield_from(5)
 
         self.loop.run_until_complete(foo())