mypy

Author: graingert

asgiref/_asyncio.py

@@ -6,23 +6,33 @@
 ]
 
 import asyncio
+import concurrent.futures
 import contextvars
 import functools
 import sys
+import types
 from asyncio import get_running_loop
-from collections.abc import Callable
-from typing import Any, TypeVar
+from collections.abc import Callable, Coroutine
+from typing import Any, Protocol, TypeVar, Union, Generic
 
 from ._context import restore_context as _restore_context
 
 _R = TypeVar("_R")
 
+Coro = Coroutine[Any, Any, _R]
 
-def create_task_threadsafe(loop, awaitable) -> None:
+
+def create_task_threadsafe(
+    loop: asyncio.AbstractEventLoop, awaitable: Coro[object]
+) -> None:
     loop.call_soon_threadsafe(loop.create_task, awaitable)
 
 
-async def wrap_task_context(loop, task_context, awaitable):
+async def wrap_task_context(
+    loop: asyncio.AbstractEventLoop,
+    task_context: list[asyncio.Task[Any]],
+    awaitable: Coro[_R],
+) -> _R:
     if task_context is None:
         return await awaitable
 
@@ -37,8 +47,30 @@ async def wrap_task_context(loop, task_context, awaitable):
         task_context.remove(current_task)
 
 
+ExcInfo = Union[
+    tuple[type[BaseException], BaseException, types.TracebackType],
+    tuple[None, None, None],
+]
+
+
+class ThreadHandlerType(Protocol, Generic[_R]):
+    def __call__(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        exc_info: ExcInfo,
+        task_context: list[asyncio.Task[Any]],
+        func: Callable[[Callable[[], _R]], _R],
+        child: Callable[[], _R],
+    ) -> _R:
+        ...
+
+
 async def run_in_executor(
-    *, loop, executor, thread_handler, child: Callable[[], _R]
+    *,
+    loop: asyncio.AbstractEventLoop,
+    executor: concurrent.futures.ThreadPoolExecutor,
+    thread_handler: ThreadHandlerType[_R],
+    child: Callable[[], _R],
 ) -> _R:
     context = contextvars.copy_context()
     func = context.run

asgiref/_trio.py

@@ -2,7 +2,10 @@
 import contextvars
 import functools
 import sys
-from typing import Any
+import types
+from collections.abc import Callable, Coroutine, Awaitable
+from typing import Any, Protocol, TypeVar, Union, Generic
+import concurrent.futures
 
 import sniffio
 import trio.lowlevel
@@ -11,12 +14,20 @@
 from . import _asyncio
 from ._context import restore_context as _restore_context
 
+_R = TypeVar("_R")
+
+Coro = Coroutine[Any, Any, _R]
+
+Loop = Union[asyncio.AbstractEventLoop, trio.lowlevel.TrioToken]
+TaskContext = list[Any]
+
 
 class TrioThreadCancelled(BaseException):
     pass
 
 
-def get_running_loop():
+def get_running_loop() -> Loop:
+
     try:
         asynclib = sniffio.current_async_library()
     except sniffio.AsyncLibraryNotFoundError:
@@ -25,16 +36,16 @@ def get_running_loop():
     if asynclib == "asyncio":
         return asyncio.get_running_loop()
     if asynclib == "trio":
-        return trio.lowlevel.current_token()
+        return trio.lowlevel.current_trio_token()
     raise RuntimeError(f"unsupported library {asynclib}")
 
 
 @trio.lowlevel.disable_ki_protection
-async def wrap_awaitable(awaitable):
+async def wrap_awaitable(awaitable: Awaitable[_R]) -> _R:
     return await awaitable
 
 
-def create_task_threadsafe(loop, awaitable):
+def create_task_threadsafe(loop: Loop, awaitable: Coro[_R]) -> None:
     if isinstance(loop, trio.lowlevel.TrioToken):
         try:
             loop.run_sync_soon(
@@ -44,15 +55,34 @@ def create_task_threadsafe(loop, awaitable):
             )
         except trio.RunFinishedError:
             raise RuntimeError("trio loop no-longer running")
+        return
+
+    _asyncio.create_task_threadsafe(loop, awaitable)
+
+
+ExcInfo = Union[
+    tuple[type[BaseException], BaseException, types.TracebackType],
+    tuple[None, None, None],
+]
+
 
-    return _asyncio.create_task_threadsafe(loop, awaitable)
+class ThreadHandlerType(Protocol, Generic[_R]):
+    def __call__(
+        self,
+        loop: Loop,
+        exc_info: ExcInfo,
+        task_context: TaskContext,
+        func: Callable[[Callable[[], _R]], _R],
+        child: Callable[[], _R],
+    ) -> _R:
+        ...
 
 
-async def run_in_executor(*, loop, executor, thread_handler, child):
+async def run_in_executor(*, loop: Loop, executor: concurrent.futures.ThreadPoolExecutor, thread_handler: ThreadHandlerType[_R], child: Callable[[], _R]) -> _R:
     if isinstance(loop, trio.lowlevel.TrioToken):
         context = contextvars.copy_context()
         func = context.run
-        task_context: list[asyncio.Task[Any]] = []
+        task_context: TaskContext = []
 
         # Run the code in the right thread
         full_func = functools.partial(
@@ -66,7 +96,7 @@ async def run_in_executor(*, loop, executor, thread_handler, child):
         try:
             if executor is None:
 
-                async def handle_cancel():
+                async def handle_cancel() -> None:
                     try:
                         await trio.sleep_forever()
                     except trio.Cancelled:
@@ -84,16 +114,17 @@ async def handle_cancel():
                         pass
                     finally:
                         nursery.cancel_scope.cancel()
+                assert False
             else:
                 event = trio.Event()
 
-                def callback(fut):
+                def callback(fut: object) -> None:
                     loop.run_sync_soon(event.set)
 
                 fut = executor.submit(full_func)
                 fut.add_done_callback(callback)
 
-                async def handle_cancel_fut():
+                async def handle_cancel_fut() -> None:
                     try:
                         await trio.sleep_forever()
                     except trio.Cancelled:
@@ -111,15 +142,17 @@ async def handle_cancel_fut():
                             return fut.result()
                         except TrioThreadCancelled:
                             pass
+                assert False
         finally:
             _restore_context(context)
 
-    return await _asyncio.run_in_executor(
-        loop=loop, executor=executor, thread_handler=thread_handler, func=func
+    else:
+        return await _asyncio.run_in_executor(
+        loop=loop, executor=executor, thread_handler=thread_handler, child=child
     )
 
 
-async def wrap_task_context(loop, task_context, awaitable):
+async def wrap_task_context(loop: Loop, task_context: Union[TaskContext, None], awaitable: Coro[_R]) -> _R:
     if task_context is None:
         return await awaitable
 
@@ -130,7 +163,6 @@ async def wrap_task_context(loop, task_context, awaitable):
                 return await awaitable
             finally:
                 task_context.remove(scope)
-        if scope.cancelled_caught:
-            raise TrioThreadCancelled
+        raise TrioThreadCancelled
 
     return await _asyncio.wrap_task_context(loop, task_context, awaitable)

asgiref/sync.py

@@ -165,7 +165,7 @@ def __init__(
         ],
         force_new_loop: Union[LoopType, bool] = False,
     ):
-        if force_new_loop and not isinstance(LoopType):
+        if force_new_loop and not isinstance(force_new_loop, LoopType):
             force_new_loop = LoopType.ASYNCIO
 
         if not callable(awaitable) or (
@@ -319,14 +319,15 @@ async def main_wrap(
         if context is not None:
             _restore_context(context[0])
 
+        result: _R
         try:
             # If we have an exception, run the function inside the except block
             # after raising it so exc_info is correctly populated.
             if exc_info[1]:
                 try:
                     raise exc_info[1]
                 except BaseException:
-                    result = await wrap_task_context(loop, task_context, awaitable)
+                    result =  await wrap_task_context(loop, task_context, awaitable)
             else:
                 result = await wrap_task_context(loop, task_context, awaitable)
         except BaseException as e:

tox.ini

@@ -11,6 +11,7 @@ commands =
     mypy: mypy . {posargs}
 deps =
     setuptools
+    mypy: trio
 
 [testenv:qa]
 skip_install = true