Support PEP 492 native coroutines.

6 files changed, 116 insertions, 27 deletions

diff --git a/asyncio/base_events.py b/asyncio/base_events.py
index 98aadaf..38344a7 100644
--- a/asyncio/base_events.py
+++ b/asyncio/base_events.py
@@ -191,8 +191,8 @@ class BaseEventLoop(events.AbstractEventLoop):
         self._thread_id = None
         self._clock_resolution = time.get_clock_info('monotonic').resolution
         self._exception_handler = None
-        self._debug = (not sys.flags.ignore_environment
-                       and bool(os.environ.get('PYTHONASYNCIODEBUG')))
+        self.set_debug((not sys.flags.ignore_environment
+                        and bool(os.environ.get('PYTHONASYNCIODEBUG'))))
         # In debug mode, if the execution of a callback or a step of a task
         # exceed this duration in seconds, the slow callback/task is logged.
         self.slow_callback_duration = 0.1
@@ -360,13 +360,18 @@ class BaseEventLoop(events.AbstractEventLoop):
             return
         if self._debug:
             logger.debug("Close %r", self)
-        self._closed = True
-        self._ready.clear()
-        self._scheduled.clear()
-        executor = self._default_executor
-        if executor is not None:
-            self._default_executor = None
-            executor.shutdown(wait=False)
+        try:
+            self._closed = True
+            self._ready.clear()
+            self._scheduled.clear()
+            executor = self._default_executor
+            if executor is not None:
+                self._default_executor = None
+                executor.shutdown(wait=False)
+        finally:
+            # It is important to unregister "sys.coroutine_wrapper"
+            # if it was registered.
+            self.set_debug(False)
 
     def is_closed(self):
         """Returns True if the event loop was closed."""
@@ -1199,3 +1204,27 @@ class BaseEventLoop(events.AbstractEventLoop):
 
     def set_debug(self, enabled):
         self._debug = enabled
+        wrapper = coroutines.debug_wrapper
+
+        try:
+            set_wrapper = sys.set_coroutine_wrapper
+        except AttributeError:
+            pass
+        else:
+            current_wrapper = sys.get_coroutine_wrapper()
+            if enabled:
+                if current_wrapper not in (None, wrapper):
+                    warnings.warn(
+                        "loop.set_debug(True): cannot set debug coroutine "
+                        "wrapper; another wrapper is already set %r" %
+                        current_wrapper, RuntimeWarning)
+                else:
+                    set_wrapper(wrapper)
+            else:
+                if current_wrapper not in (None, wrapper):
+                    warnings.warn(
+                        "loop.set_debug(False): cannot unset debug coroutine "
+                        "wrapper; another wrapper was set %r" %
+                        current_wrapper, RuntimeWarning)
+                else:
+                    set_wrapper(None)
diff --git a/asyncio/coroutines.py b/asyncio/coroutines.py
index c639461..20c4579 100644
--- a/asyncio/coroutines.py
+++ b/asyncio/coroutines.py
@@ -14,6 +14,9 @@ from . import futures
 from .log import logger
 
 
+_PY35 = sys.version_info >= (3, 5)
+
+
 # Opcode of "yield from" instruction
 _YIELD_FROM = opcode.opmap['YIELD_FROM']
 
@@ -30,6 +33,27 @@ _DEBUG = (not sys.flags.ignore_environment
           and bool(os.environ.get('PYTHONASYNCIODEBUG')))
 
 
+try:
+    types.coroutine
+except AttributeError:
+    native_coroutine_support = False
+else:
+    native_coroutine_support = True
+
+try:
+    _iscoroutinefunction = inspect.iscoroutinefunction
+except AttributeError:
+    _iscoroutinefunction = lambda func: False
+
+try:
+    inspect.CO_COROUTINE
+except AttributeError:
+    _is_native_coro_code = lambda code: False
+else:
+    _is_native_coro_code = lambda code: (code.co_flags &
+                                         inspect.CO_COROUTINE)
+
+
 # Check for CPython issue #21209
 def has_yield_from_bug():
     class MyGen:
@@ -54,16 +78,27 @@ _YIELD_FROM_BUG = has_yield_from_bug()
 del has_yield_from_bug
 
 
+def debug_wrapper(gen):
+    # This function is called from 'sys.set_coroutine_wrapper'.
+    # We only wrap here coroutines defined via 'async def' syntax.
+    # Generator-based coroutines are wrapped in @coroutine
+    # decorator.
+    if _is_native_coro_code(gen.gi_code):
+        return CoroWrapper(gen, None)
+    else:
+        return gen
+
+
 class CoroWrapper:
     # Wrapper for coroutine object in _DEBUG mode.
 
-    def __init__(self, gen, func):
-        assert inspect.isgenerator(gen), gen
+    def __init__(self, gen, func=None):
+        assert inspect.isgenerator(gen) or inspect.iscoroutine(gen), gen
         self.gen = gen
-        self.func = func
+        self.func = func # Used to unwrap @coroutine decorator
         self._source_traceback = traceback.extract_stack(sys._getframe(1))
-        # __name__, __qualname__, __doc__ attributes are set by the coroutine()
-        # decorator
+        self.__name__ = getattr(gen, '__name__', None)
+        self.__qualname__ = getattr(gen, '__qualname__', None)
 
     def __repr__(self):
         coro_repr = _format_coroutine(self)
@@ -75,6 +110,9 @@ class CoroWrapper:
     def __iter__(self):
         return self
 
+    if _PY35:
+        __await__ = __iter__ # make compatible with 'await' expression
+
     def __next__(self):
         return next(self.gen)
 
@@ -133,6 +171,14 @@ def coroutine(func):
     If the coroutine is not yielded from before it is destroyed,
     an error message is logged.
     """
+    is_coroutine = _iscoroutinefunction(func)
+    if is_coroutine and _is_native_coro_code(func.__code__):
+        # In Python 3.5 that's all we need to do for coroutines
+        # defiend with "async def".
+        # Wrapping in CoroWrapper will happen via
+        # 'sys.set_coroutine_wrapper' function.
+        return func
+
     if inspect.isgeneratorfunction(func):
         coro = func
     else:
@@ -144,18 +190,22 @@ def coroutine(func):
             return res
 
     if not _DEBUG:
-        wrapper = coro
+        if native_coroutine_support:
+            wrapper = types.coroutine(coro)
+        else:
+            wrapper = coro
     else:
         @functools.wraps(func)
         def wrapper(*args, **kwds):
-            w = CoroWrapper(coro(*args, **kwds), func)
+            w = CoroWrapper(coro(*args, **kwds), func=func)
             if w._source_traceback:
                 del w._source_traceback[-1]
-            if hasattr(func, '__name__'):
-                w.__name__ = func.__name__
-            if hasattr(func, '__qualname__'):
-                w.__qualname__ = func.__qualname__
-            w.__doc__ = func.__doc__
+            # Python < 3.5 does not implement __qualname__
+            # on generator objects, so we set it manually.
+            # We use getattr as some callables (such as
+            # functools.partial may lack __qualname__).
+            w.__name__ = getattr(func, '__name__', None)
+            w.__qualname__ = getattr(func, '__qualname__', None)
             return w
 
     wrapper._is_coroutine = True  # For iscoroutinefunction().
@@ -164,7 +214,8 @@ def coroutine(func):
 
 def iscoroutinefunction(func):
     """Return True if func is a decorated coroutine function."""
-    return getattr(func, '_is_coroutine', False)
+    return (getattr(func, '_is_coroutine', False) or
+            _iscoroutinefunction(func))
 
 
 _COROUTINE_TYPES = (types.GeneratorType, CoroWrapper)
diff --git a/asyncio/futures.py b/asyncio/futures.py
index 74a99ba..d06828a 100644
--- a/asyncio/futures.py
+++ b/asyncio/futures.py
@@ -19,6 +19,7 @@ _CANCELLED = 'CANCELLED'
 _FINISHED = 'FINISHED'
 
 _PY34 = sys.version_info >= (3, 4)
+_PY35 = sys.version_info >= (3, 5)
 
 Error = concurrent.futures._base.Error
 CancelledError = concurrent.futures.CancelledError
@@ -387,6 +388,9 @@ class Future:
         assert self.done(), "yield from wasn't used with future"
         return self.result()  # May raise too.
 
+    if _PY35:
+        __await__ = __iter__ # make compatible with 'await' expression
+
 
 def wrap_future(fut, *, loop=None):
     """Wrap concurrent.futures.Future object."""
diff --git a/asyncio/tasks.py b/asyncio/tasks.py
index f617b62..fcb3833 100644
--- a/asyncio/tasks.py
+++ b/asyncio/tasks.py
@@ -11,6 +11,7 @@ import functools
 import inspect
 import linecache
 import sys
+import types
 import traceback
 import warnings
 import weakref
@@ -73,7 +74,10 @@ class Task(futures.Future):
         super().__init__(loop=loop)
         if self._source_traceback:
             del self._source_traceback[-1]
-        self._coro = iter(coro)  # Use the iterator just in case.
+        if coro.__class__ is types.GeneratorType:
+            self._coro = coro
+        else:
+            self._coro = iter(coro)  # Use the iterator just in case.
         self._fut_waiter = None
         self._must_cancel = False
         self._loop.call_soon(self._step)
@@ -236,7 +240,7 @@ class Task(futures.Future):
             elif value is not None:
                 result = coro.send(value)
             else:
-                result = next(coro)
+                result = coro.send(None)
         except StopIteration as exc:
             self.set_result(exc.value)
         except futures.CancelledError as exc:
diff --git a/tests/test_base_events.py b/tests/test_base_events.py
index 8c4498c..b1f1e56 100644
--- a/tests/test_base_events.py
+++ b/tests/test_base_events.py
@@ -61,7 +61,8 @@ class BaseEventLoopTests(test_utils.TestCase):
             NotImplementedError,
             self.loop._make_write_pipe_transport, m, m)
         gen = self.loop._make_subprocess_transport(m, m, m, m, m, m, m)
-        self.assertRaises(NotImplementedError, next, iter(gen))
+        with self.assertRaises(NotImplementedError):
+            gen.send(None)
 
     def test_close(self):
         self.assertFalse(self.loop.is_closed())
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index 4119085..6541df7 100644
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -1638,7 +1638,7 @@ class TaskTests(test_utils.TestCase):
             return a
 
         def call(arg):
-            cw = asyncio.coroutines.CoroWrapper(foo(), foo)
+            cw = asyncio.coroutines.CoroWrapper(foo())
             cw.send(None)
             try:
                 cw.send(arg)
@@ -1653,7 +1653,7 @@ class TaskTests(test_utils.TestCase):
     def test_corowrapper_weakref(self):
         wd = weakref.WeakValueDictionary()
         def foo(): yield from []
-        cw = asyncio.coroutines.CoroWrapper(foo(), foo)
+        cw = asyncio.coroutines.CoroWrapper(foo())
         wd['cw'] = cw  # Would fail without __weakref__ slot.
         cw.gen = None  # Suppress warning from __del__.