split out trio/asyncio compat

Author: graingert

asgiref/_asyncio.py

@@ -0,0 +1,82 @@
+__all__ = [
+    "get_running_loop",
+    "create_task_threadsafe",
+    "wrap_task_context",
+    "run_in_executor",
+]
+
+import asyncio
+import contextvars
+import functools
+import sys
+from asyncio import get_running_loop
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from ._context import restore_context as _restore_context
+
+_R = TypeVar("_R")
+
+
+def create_task_threadsafe(loop, awaitable) -> None:
+    loop.call_soon_threadsafe(loop.create_task, awaitable)
+
+
+async def wrap_task_context(loop, task_context, awaitable):
+    if task_context is None:
+        return await awaitable
+
+    current_task = asyncio.current_task(loop)
+    if current_task is None:
+        return await awaitable
+
+    task_context.append(current_task)
+    try:
+        return await awaitable
+    finally:
+        task_context.remove(current_task)
+
+
+async def run_in_executor(
+    *, loop, executor, thread_handler, child: Callable[[], _R]
+) -> _R:
+    context = contextvars.copy_context()
+    func = context.run
+    task_context: list[asyncio.Task[Any]] = []
+
+    # Run the code in the right thread
+    exec_coro = loop.run_in_executor(
+        executor,
+        functools.partial(
+            thread_handler,
+            loop,
+            sys.exc_info(),
+            task_context,
+            func,
+            child,
+        ),
+    )
+    ret: _R
+    try:
+        ret = await asyncio.shield(exec_coro)
+    except asyncio.CancelledError:
+        cancel_parent = True
+        try:
+            task = task_context[0]
+            task.cancel()
+            try:
+                await task
+                cancel_parent = False
+            except asyncio.CancelledError:
+                pass
+        except IndexError:
+            pass
+        if exec_coro.done():
+            raise
+        if cancel_parent:
+            exec_coro.cancel()
+        ret = await exec_coro
+    finally:
+        _restore_context(context)
+
+    return ret

asgiref/_context.py

@@ -0,0 +1,13 @@
+import contextvars
+
+
+def restore_context(context: contextvars.Context) -> None:
+    # Check for changes in contextvars, and set them to the current
+    # context for downstream consumers
+    for cvar in context:
+        cvalue = context.get(cvar)
+        try:
+            if cvar.get() != cvalue:
+                cvar.set(cvalue)
+        except LookupError:
+            cvar.set(cvalue)

asgiref/_trio.py

@@ -0,0 +1,136 @@
+import asyncio
+import contextvars
+import functools
+import sys
+from typing import Any
+
+import sniffio
+import trio.lowlevel
+import trio.to_thread
+
+from . import _asyncio
+from ._context import restore_context as _restore_context
+
+
+class TrioThreadCancelled(BaseException):
+    pass
+
+
+def get_running_loop():
+    try:
+        asynclib = sniffio.current_async_library()
+    except sniffio.AsyncLibraryNotFoundError:
+        return asyncio.get_running_loop()
+
+    if asynclib == "asyncio":
+        return asyncio.get_running_loop()
+    if asynclib == "trio":
+        return trio.lowlevel.current_token()
+    raise RuntimeError(f"unsupported library {asynclib}")
+
+
+@trio.lowlevel.disable_ki_protection
+async def wrap_awaitable(awaitable):
+    return await awaitable
+
+
+def create_task_threadsafe(loop, awaitable):
+    if isinstance(loop, trio.lowlevel.TrioToken):
+        try:
+            loop.run_sync_soon(
+                trio.lowlevel.spawn_system_task,
+                wrap_awaitable,
+                awaitable,
+            )
+        except trio.RunFinishedError:
+            raise RuntimeError("trio loop no-longer running")
+
+    return _asyncio.create_task_threadsafe(loop, awaitable)
+
+
+async def run_in_executor(*, loop, executor, thread_handler, child):
+    if isinstance(loop, trio.lowlevel.TrioToken):
+        context = contextvars.copy_context()
+        func = context.run
+        task_context: list[asyncio.Task[Any]] = []
+
+        # Run the code in the right thread
+        full_func = functools.partial(
+            thread_handler,
+            loop,
+            sys.exc_info(),
+            task_context,
+            func,
+            child,
+        )
+        try:
+            if executor is None:
+
+                async def handle_cancel():
+                    try:
+                        await trio.sleep_forever()
+                    except trio.Cancelled:
+                        if task_context:
+                            task_context[0].cancel()
+                        raise
+
+                async with trio.open_nursery() as nursery:
+                    nursery.start_soon(handle_cancel)
+                    try:
+                        return await trio.to_thread.run_sync(
+                            thread_handler, func, abandon_on_cancel=False
+                        )
+                    except TrioThreadCancelled:
+                        pass
+                    finally:
+                        nursery.cancel_scope.cancel()
+            else:
+                event = trio.Event()
+
+                def callback(fut):
+                    loop.run_sync_soon(event.set)
+
+                fut = executor.submit(full_func)
+                fut.add_done_callback(callback)
+
+                async def handle_cancel_fut():
+                    try:
+                        await trio.sleep_forever()
+                    except trio.Cancelled:
+                        fut.cancel()
+                        if task_context:
+                            task_context[0].cancel()
+                        raise
+
+                async with trio.open_nursery() as nursery:
+                    nursery.start_soon(handle_cancel_fut)
+                    with trio.CancelScope(shield=True):
+                        await event.wait()
+                        nursery.cancel_scope.cancel()
+                        try:
+                            return fut.result()
+                        except TrioThreadCancelled:
+                            pass
+        finally:
+            _restore_context(context)
+
+    return await _asyncio.run_in_executor(
+        loop=loop, executor=executor, thread_handler=thread_handler, func=func
+    )
+
+
+async def wrap_task_context(loop, task_context, awaitable):
+    if task_context is None:
+        return await awaitable
+
+    if isinstance(loop, trio.lowlevel.TrioToken):
+        with trio.CancelScope() as scope:
+            task_context.append(scope)
+            try:
+                return await awaitable
+            finally:
+                task_context.remove(scope)
+        if scope.cancelled_caught:
+            raise TrioThreadCancelled
+
+    return await _asyncio.wrap_task_context(loop, task_context, awaitable)