Enable async state/message handlers (asyncio/trio/curio) with ContextVar preservation and tests

Author: Nikhil172913832

src/pact/_util.py

@@ -9,6 +9,8 @@
 
 from __future__ import annotations
 
+import asyncio
+import contextvars
 import inspect
 import logging
 import socket
@@ -18,6 +20,23 @@
 from inspect import Parameter, _ParameterKind
 from typing import TYPE_CHECKING, TypeVar
 
+try:
+    import sniffio
+except ImportError:
+    sniffio = None  # type: ignore[assignment]
+
+try:
+    import trio
+    from trio.lowlevel import current_trio_token
+except ImportError:
+    trio = None  # type: ignore[assignment]
+    current_trio_token = None  # type: ignore[assignment]
+
+try:
+    import curio
+except (ImportError, AttributeError):
+    curio = None  # type: ignore[assignment]
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
 
@@ -179,7 +198,7 @@ def find_free_port() -> int:
         return s.getsockname()[1]
 
 
-def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa: C901
+def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa: C901, PLR0912
     """
     Apply arguments to a function.
 
@@ -188,6 +207,9 @@ def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa:
     it is possible to pass arguments by name, and falling back to positional
     arguments if not.
 
+    This function supports both synchronous and asynchronous callables. If the
+    callable is a coroutine function, it will be executed in an event loop.
+
     Args:
         f:
             The function to apply the arguments to.
@@ -200,6 +222,7 @@ def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa:
     Returns:
         The result of the function.
     """
+    is_async = inspect.iscoroutinefunction(f)
     signature = inspect.signature(f)
     f_name = (
         f.__qualname__
@@ -226,7 +249,15 @@ def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa:
     # First, we inspect the keyword arguments and try and pass in some arguments
     # by currying them in.
     for param in signature.parameters.values():
-        if param.name not in args:
+        # Try matching the parameter name, or if it starts with underscore,
+        # also try matching without the leading underscore.
+        arg_key = None
+        if param.name in args:
+            arg_key = param.name
+        elif param.name.startswith("_") and param.name[1:] in args:
+            arg_key = param.name[1:]
+
+        if arg_key is None:
             # If a parameter is not known, we will ignore it.
             #
             # If the ignored parameter doesn't have a default value, it will
@@ -246,12 +277,12 @@ def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa:
         if param.kind in positional_match:
             # We iterate through the parameters in order that they are defined,
             # making it fine to pass them in by position one at a time.
-            f = partial(f, args[param.name])
-            del args[param.name]
+            f = partial(f, args[arg_key])
+            del args[arg_key]
 
         if param.kind in keyword_match:
-            f = partial(f, **{param.name: args[param.name]})
-            del args[param.name]
+            f = partial(f, **{param.name: args[arg_key]})
+            del args[arg_key]
             continue
 
     # At this stage, we have checked all arguments. If we have any arguments
@@ -282,7 +313,107 @@ def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa:
             )
 
     try:
+        if is_async:
+            return _run_async(f)
         return f()
     except Exception:
         logger.exception("Error occurred while calling function %s", f_name)
         raise
+
+
+def _run_async(coro_func: Callable[..., _T]) -> _T:
+    """
+    Run an async function in an event loop.
+
+    Detects the async runtime (asyncio, trio, or curio) and executes the
+    coroutine appropriately. Preserves ContextVars when creating a new event
+    loop, which is important when handlers are called from threads.
+
+    Args:
+        coro_func:
+            The async function to run (already partially applied).
+
+    Returns:
+        The result of the async function.
+
+    Raises:
+        RuntimeError:
+            If the detected runtime is not installed.
+    """
+    runtime = _detect_async_runtime_from_coro(coro_func)
+
+    if runtime == "trio":
+        if trio is None:
+            msg = "trio is not installed"
+            raise RuntimeError(msg)
+
+        if current_trio_token is not None:
+            try:
+                token = current_trio_token()
+                return trio.from_thread.run_sync(coro_func, trio_token=token)
+            except RuntimeError:
+                pass
+
+        ctx = contextvars.copy_context()
+        return ctx.run(trio.run, coro_func)
+
+    if runtime == "curio":
+        if curio is None:
+            msg = "curio is not installed"
+            raise RuntimeError(msg)
+
+        try:
+            return curio.AWAIT(coro_func())
+        except RuntimeError:
+            ctx = contextvars.copy_context()
+            return ctx.run(curio.run, coro_func)
+
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop is not None:
+        future = asyncio.run_coroutine_threadsafe(coro_func(), loop)
+        return future.result()
+
+    ctx = contextvars.copy_context()
+    return ctx.run(asyncio.run, coro_func())
+
+
+def _detect_async_runtime_from_coro(coro_func: Callable[..., _T]) -> str:
+    """
+    Detect async runtime by inspecting the coroutine function.
+
+    Args:
+        coro_func:
+            The async function to inspect. May be a partial object.
+
+    Returns:
+        The detected runtime: "asyncio", "trio", or "curio".
+    """
+    if sniffio is not None:
+        try:
+            return sniffio.current_async_library()
+        except sniffio.AsyncLibraryNotFoundError:
+            pass
+
+    func = getattr(coro_func, "func", coro_func)
+    func_code = getattr(func, "__code__", None)
+    func_globals = getattr(func, "__globals__", {})
+
+    if func_code is not None:
+        code_names = set(func_code.co_names)
+
+        trio_indicators = {"trio", "open_nursery", "current_trio_token"}
+        curio_indicators = {"curio", "TaskGroup", "AWAIT"}
+
+        uses_trio = bool(trio_indicators & code_names)
+        uses_curio = bool(curio_indicators & code_names)
+
+        if uses_trio and "trio" in func_globals:
+            return "trio"
+        if uses_curio and "curio" in func_globals:
+            return "curio"
+
+    return "asyncio"