loop: Set coroutine wrapper in debug mode

Author: 1st1

uvloop/includes/stdlib.pxi

@@ -1,4 +1,5 @@
-import asyncio, asyncio.log, asyncio.base_events, asyncio.sslproto
+import asyncio, asyncio.log, asyncio.base_events, \
+       asyncio.sslproto, asyncio.coroutines
 import collections
 import concurrent.futures
 import functools
@@ -30,6 +31,7 @@ cdef aio_wrap_future = asyncio.wrap_future
 cdef aio_BaseProtocol = asyncio.BaseProtocol
 cdef aio_Protocol = asyncio.Protocol
 cdef aio_SSLProtocol = asyncio.sslproto.SSLProtocol
+cdef aio_debug_wrapper = asyncio.coroutines.debug_wrapper
 
 cdef col_deque = collections.deque
 cdef col_Iterable = collections.Iterable
@@ -78,6 +80,8 @@ cdef os_O_RDWR = os.O_RDWR
 
 cdef sys_ignore_environment = sys.flags.ignore_environment
 cdef sys_exc_info = sys.exc_info
+cdef sys_set_coroutine_wrapper = sys.set_coroutine_wrapper
+cdef sys_get_coroutine_wrapper = sys.get_coroutine_wrapper
 
 cdef ssl_SSLContext = ssl.SSLContext
 

uvloop/loop.pxd

@@ -35,6 +35,8 @@ cdef class Loop:
     cdef:
         uv.uv_loop_t *uvloop
 
+        bint _coroutine_wrapper_set
+
         readonly bint _closed
         bint _debug
         bint _running
@@ -156,6 +158,8 @@ cdef class Loop:
     cdef _sock_connect(self, fut, sock, address)
     cdef _sock_connect_cb(self, fut, sock, address)
 
+    cdef _set_coroutine_wrapper(self, bint enabled)
+
 
 include "cbhandles.pxd"
 

uvloop/loop.pyx

@@ -149,6 +149,8 @@ cdef class Loop:
 
         uv.uv_disable_stdio_inheritance()
 
+        self._coroutine_wrapper_set = False
+
     def __init__(self):
         self.set_debug((not sys_ignore_environment
                         and bool(os_environ.get('PYTHONASYNCIODEBUG'))))
@@ -684,6 +686,33 @@ cdef class Loop:
         else:
             fut.set_result(None)
 
+    cdef _set_coroutine_wrapper(self, bint enabled):
+        enabled = bool(enabled)
+        if self._coroutine_wrapper_set == enabled:
+            return
+
+        wrapper = aio_debug_wrapper
+        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:
+                sys_set_coroutine_wrapper(wrapper)
+                self._coroutine_wrapper_set = True
+        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:
+                sys_set_coroutine_wrapper(None)
+                self._coroutine_wrapper_set = False
+
     IF DEBUG:
         def print_debug_info(self):
             cdef:
@@ -833,7 +862,11 @@ cdef class Loop:
             # loop.stop() was called right before loop.run_forever().
             # This is how asyncio loop behaves.
             mode = uv.UV_RUN_NOWAIT
-        self._run(mode)
+        self._set_coroutine_wrapper(self._debug)
+        try:
+            self._run(mode)
+        finally:
+            self._set_coroutine_wrapper(0)
 
     def close(self):
         self._close()
@@ -850,6 +883,9 @@ cdef class Loop:
         else:
             self._debug = 0
 
+        if self.is_running():
+            self._set_coroutine_wrapper(bool(enabled))
+
     def is_running(self):
         if self._running == 0:
             return False