Stricter handling in Caller class.

Author: Alan Fleming

src/async_kernel/caller.py

@@ -371,9 +371,9 @@ class Caller:
     _backend: Backend
     _queue_map: weakref.WeakKeyDictionary[Callable[..., Awaitable[Any]], MemoryObjectSendStream[tuple]]
     _taskgroup: TaskGroup | None = None
-    _callers: deque[tuple[contextvars.Context, Future] | Callable[[], Any]]
+    _jobs: deque[tuple[contextvars.Context, Future] | Callable[[], Any]]
     _thread: threading.Thread
-    _callers_added: threading.Event
+    _job_added: threading.Event
     _stopped_event: threading.Event
     _stopped = False
     _protected = False
@@ -428,8 +428,8 @@ def __new__(
             inst._backend = Backend(sniffio.current_async_library())
             inst._thread = thread
             inst.log = log or logging.LoggerAdapter(logging.getLogger())
-            inst._callers = deque()
-            inst._callers_added = threading.Event()
+            inst._jobs = deque()
+            inst._job_added = threading.Event()
             inst._protected = protected
             inst._queue_map = weakref.WeakKeyDictionary()
             cls._instances[thread] = inst
@@ -461,24 +461,27 @@ async def _server_loop(self, tg: TaskGroup, task_status: TaskStatus[None]) -> No
             self.iopub_sockets[self.thread] = socket
             task_status.started()
             while not self._stopped:
-                if not self._callers:
-                    self._callers_added.clear()
-                await wait_thread_event(self._callers_added)
-                while self._callers:
+                if not self._jobs:
+                    self._job_added.clear()
+                await wait_thread_event(self._job_added)
+                while self._jobs:
                     if self._stopped:
                         return
-                    job = self._callers.popleft()
+                    job = self._jobs.popleft()
                     if isinstance(job, Callable):
                         try:
-                            job()
+                            if inspect.iscoroutinefunction(job):
+                                await job()
+                            else:
+                                job()
                         except Exception as e:
                             self.log.exception("Simple call failed", exc_info=e)
                     else:
                         context, fut = job
                         context.run(tg.start_soon, self._wrap_call, fut)
         finally:
             self._running = False
-            for job in self._callers:
+            for job in self._jobs:
                 if isinstance(job, tuple):
                     job[1].set_exception(FutureCancelledError())
             socket.close()
@@ -493,8 +496,8 @@ def _schedule_wrapped_call(self, func: Callable, /, args: tuple, kwargs: dict, *
         if threading.current_thread() is self.thread and (tg := self._taskgroup):
             tg.start_soon(self._wrap_call, fut)
         else:
-            self._callers.append((contextvars.copy_context(), fut))
-            self._callers_added.set()
+            self._jobs.append((contextvars.copy_context(), fut))
+            self._job_added.set()
         return fut
 
     async def _wrap_call(self, fut: Future) -> None:
@@ -510,9 +513,12 @@ async def _wrap_call(self, fut: Future) -> None:
                 try:
                     if (delay := md.get("delay")) and ((delay := delay - time.monotonic() + md["start_time"]) > 0):
                         await anyio.sleep(delay)
-                    result = func(*md["args"], **md["kwargs"]) if callable(func) else func
-                    if inspect.isawaitable(result) and result is not fut:
-                        result = await result
+                    # Evaluate
+                    if inspect.iscoroutinefunction(func):
+                        result = await func(*md["args"], **md["kwargs"])
+                    else:
+                        result = func(*md["args"], **md["kwargs"])
+                    # Cancellation
                     if fut.cancelled() and not scope.cancel_called:
                         scope.cancel()
                     fut.set_result(result)
@@ -566,15 +572,20 @@ def stop(self, *, force=False) -> None:
         for sender in self._queue_map.values():
             sender.close()
         self._queue_map.clear()
-        self._callers_added.set()
+        self._job_added.set()
         self._instances.pop(self.thread, None)
         if self in self._to_thread_pool:
             self._to_thread_pool.remove(self)
         if self.thread is not threading.current_thread():
             self._stopped_event.wait()
 
     def call_later(
-        self, delay: float, func: Callable[P, T | Awaitable[T]], /, *args: P.args, **kwargs: P.kwargs
+        self,
+        delay: float,
+        func: Callable[P, T | CoroutineType[Any, Any, T]],
+        /,
+        *args: P.args,
+        **kwargs: P.kwargs,
     ) -> Future[T]:
         """
         Schedule func to be called in caller's event loop copying the current context.
@@ -587,7 +598,13 @@ def call_later(
         """
         return self._schedule_wrapped_call(func, args, kwargs, delay=delay, start_time=time.monotonic())
 
-    def call_soon(self, func: Callable[P, T | Awaitable[T]], /, *args: P.args, **kwargs: P.kwargs) -> Future[T]:
+    def call_soon(
+        self,
+        func: Callable[P, T | CoroutineType[Any, Any, T]],
+        /,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Future[T]:
         """
         Schedule func to be called in caller's event loop copying the current context.
 
@@ -598,12 +615,18 @@ def call_soon(self, func: Callable[P, T | Awaitable[T]], /, *args: P.args, **kwa
         """
         return self._schedule_wrapped_call(func, args, kwargs)
 
-    def call_direct(self, func: Callable[P, Any], /, *args: P.args, **kwargs: P.kwargs) -> None:
+    def call_direct(
+        self,
+        func: Callable[P, T | CoroutineType[Any, Any, T]],
+        /,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> None:
         """
         Schedule `func` to be called in caller's event loop directly.
 
         This method is provided to facilitate lightweight *thread-safe* function calls that
-        need to be done from within the callers event loop.
+        need to be performed from within the callers event loop/taskgroup.
 
         Args:
             func: The function (awaitables permitted, though discouraged).
@@ -612,11 +635,11 @@ def call_direct(self, func: Callable[P, Any], /, *args: P.args, **kwargs: P.kwar
 
         ??? warning
 
-            - Use this method for lightweight calls only.
-            - Corroutines will **not** be awaited.
+            **Use this method for lightweight calls only!**
+
         """
-        self._callers.append(functools.partial(func, *args, **kwargs))
-        self._callers_added.set()
+        self._jobs.append(functools.partial(func, *args, **kwargs))
+        self._job_added.set()
 
     def queue_exists(self, func: Callable) -> bool:
         "Returns True if an execution queue exists for `func`."
@@ -736,13 +759,24 @@ def get_instance(cls, name: str | None = "MainThread", *, create: bool = False)
         raise RuntimeError(msg)
 
     @classmethod
-    def to_thread(cls, func: Callable[P, T | Awaitable[T]], /, *args: P.args, **kwargs: P.kwargs) -> Future[T]:
+    def to_thread(
+        cls,
+        func: Callable[P, T | CoroutineType[Any, Any, T]],
+        /,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Future[T]:
         """A classmethod to call func in a separate thread see also [to_thread_by_name][async_kernel.Caller.to_thread_by_name]."""
         return cls.to_thread_by_name(None, func, *args, **kwargs)
 
     @classmethod
     def to_thread_by_name(
-        cls, name: str | None, func: Callable[P, T | Awaitable[T]], /, *args: P.args, **kwargs: P.kwargs
+        cls,
+        name: str | None,
+        func: Callable[P, T | CoroutineType[Any, Any, T]],
+        /,
+        *args: P.args,
+        **kwargs: P.kwargs,
     ) -> Future[T]:
         """
         A classmethod to call func in the thread specified by name.

tests/test_caller.py

@@ -4,7 +4,6 @@
 import threading
 import time
 import weakref
-from contextvars import ContextVar
 from random import random
 from typing import Literal, cast
 
@@ -253,6 +252,16 @@ async def runner():
             await anyio.to_thread.run_sync(_in_thread)
         assert caller not in Caller.all_callers()
 
+    async def test_direct_async(self, caller: Caller):
+        event = AsyncEvent()
+
+        async def set_event():
+            event.set()
+
+        caller.call_direct(set_event)
+        with anyio.fail_after(1):
+            await event.wait()
+
     async def test_cancels_on_exit(self):
         is_cancelled = False
         async with Caller(create=True) as caller:
@@ -457,27 +466,6 @@ async def test_call_early(self, anyio_backend) -> None:
         async with caller:
             await fut
 
-    async def test_call_coroutine(self, caller: Caller):
-        # Test we can await a coroutine, note that it is not permitted with the type hints,
-        # but should probably be discouraged anyway since there is no way of knowing
-        # (with type hints) if a coroutine has already been awaited.
-        my_contextvar = ContextVar[int]("my_contextvar")
-        my_contextvar.set(1)
-
-        async def my_func():
-            await anyio.sleep(0)
-            assert my_contextvar.get() == 1
-            return True
-
-        # Discouraged
-        fut = Caller.to_thread(my_func())  # pyright: ignore[reportCallIssue, reportArgumentType]
-        val = await fut
-        assert val is True
-        # This the preferred way of calling.
-        fut = Caller.to_thread(my_func)
-        val = await fut
-        assert val is True
-
     async def test_current_future(self, anyio_backend):
         async with Caller(create=True) as caller:
             fut = caller.call_soon(Caller.current_future)