re-set context

Author: smurfix

trio_asyncio/base.py

@@ -214,8 +214,12 @@ async def run_coroutine(self, coro):
         This is a Trio coroutine.
         """
         self._check_closed()
+        t = sniffio.current_async_library_cvar.set("asyncio")
         coro = asyncio.ensure_future(coro, loop=self)
-        return await run_future(coro)
+        try:
+            return await run_future(coro)
+        finally:
+            sniffio.current_async_library_cvar.reset(t)
 
     def wrap_generator(self, gen, *args):
         """Call an asyncio generator.

trio_asyncio/handles.py

@@ -1,6 +1,7 @@
 import sys
 import trio
 import asyncio
+import sniffio
 try:
     from asyncio.format_helpers import _format_callback, _get_function_source
 except ImportError:  # <3.7
@@ -17,6 +18,11 @@ def _format_callback_source(func, args):
     return func_repr
 
 
+async def _set_sniff(proc, *args):
+    sniffio.current_async_library_cvar.set("trio")
+    return await proc(*args)
+
+
 class _TrioHandle:
     """
     This extends asyncio.Handle by providing:
@@ -98,9 +104,9 @@ async def _call_async(self, task_status=trio.TASK_STATUS_IGNORED):
                 with trio.open_cancel_scope() as scope:
                     self._scope = scope
                     if self._is_sync is None:
-                        await self._context.run(self._callback, self)
+                        await self._context.run(_set_sniff, self._callback, self)
                     else:
-                        await self._context.run(self._callback, *self._args)
+                        await self._context.run(_set_sniff, self._callback, *self._args)
             except Exception as exc:
                 self._raise(exc)
             finally: