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

Author: Nikhil172913832

src/pact/_util.py

@@ -9,17 +9,39 @@
 
 from __future__ import annotations
 
+import asyncio
+import contextvars
 import inspect
 import logging
 import socket
 import warnings
 from contextlib import closing
 from functools import partial
 from inspect import Parameter, _ParameterKind
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Mapping
+    from collections.abc import Callable, Coroutine, Mapping
+
+try:
+    import sniffio  # type: ignore[import-not-found]
+except ImportError:
+    sniffio = None  # type: ignore[assignment]
+
+try:
+    import trio  # type: ignore[import-not-found]
+    from trio.lowlevel import current_trio_token  # type: ignore[import-not-found]
+except ImportError:
+    trio = None  # type: ignore[assignment]
+    current_trio_token = None  # type: ignore[assignment]
+
+try:
+    import curio  # type: ignore[import-not-found,import-untyped]
+except (ImportError, AttributeError):
+    curio = None  # type: ignore[assignment]
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine, Mapping
 
 logger = logging.getLogger(__name__)
 
@@ -179,7 +201,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 +210,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 +225,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 +252,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 +280,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 +316,122 @@ def apply_args(f: Callable[..., _T], args: Mapping[str, object]) -> _T:  # noqa:
             )
 
     try:
+        if is_async:
+            result = f()
+            if inspect.iscoroutine(result):
+                return _run_async_coroutine(result)
+            return result
         return f()
     except Exception:
         logger.exception("Error occurred while calling function %s", f_name)
         raise
+
+
+def _run_async_coroutine(coro: Coroutine[Any, Any, _T]) -> _T:  # noqa: C901
+    """
+    Run a coroutine 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:
+            The coroutine to run.
+
+    Returns:
+        The result of the coroutine.
+
+    Raises:
+        RuntimeError:
+            If the detected runtime is not installed.
+    """
+    runtime = _detect_async_runtime_from_coroutine(coro)
+
+    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()
+
+                async def _run_with_token() -> _T:
+                    return await coro
+
+                return trio.from_thread.run_sync(_run_with_token, trio_token=token)  # type: ignore[return-value]
+            except RuntimeError:
+                pass
+
+        ctx = contextvars.copy_context()
+
+        async def _run_trio() -> _T:
+            return await coro
+
+        return ctx.run(trio.run, _run_trio)
+
+    if runtime == "curio":
+        if curio is None:
+            msg = "curio is not installed"
+            raise RuntimeError(msg)
+
+        try:
+            return curio.AWAIT(coro)
+        except RuntimeError:
+            ctx = contextvars.copy_context()
+
+            async def _run_curio() -> _T:
+                return await coro
+
+            return ctx.run(curio.run, _run_curio)
+
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop is not None:
+        future: asyncio.Future[_T] = asyncio.run_coroutine_threadsafe(coro, loop)  # type: ignore[arg-type,assignment]
+        return future.result()
+
+    ctx = contextvars.copy_context()
+    return ctx.run(asyncio.run, coro)  # type: ignore[arg-type,return-value]
+
+
+def _detect_async_runtime_from_coroutine(coro: Coroutine[Any, Any, _T]) -> str:
+    """
+    Detect async runtime by inspecting the coroutine object.
+
+    Args:
+        coro:
+            The coroutine object to inspect.
+
+    Returns:
+        The detected runtime: "asyncio", "trio", or "curio".
+    """
+    if sniffio is not None:
+        try:
+            return sniffio.current_async_library()
+        except sniffio.AsyncLibraryNotFoundError:
+            pass
+
+    # Inspect the coroutine's code and globals to determine runtime
+    func_code = coro.cr_code  # type: ignore[attr-defined]
+    code_names = set(func_code.co_names)
+
+    # Check if trio or curio modules are actually used in the code
+    # If the module name is in co_names, it's being accessed (e.g., trio.sleep)
+    # We also verify with specific indicators to avoid false positives
+    if "trio" in code_names:
+        trio_indicators = {"open_nursery", "current_trio_token", "CancelScope", "sleep"}
+        if any(indicator in code_names for indicator in trio_indicators):
+            return "trio"
+
+    if "curio" in code_names:
+        curio_indicators = {"TaskGroup", "spawn", "AWAIT", "sleep"}
+        if any(indicator in code_names for indicator in curio_indicators):
+            return "curio"
+
+    # Default to asyncio as it's the most common
+    return "asyncio"