Merge branch 'main' into pipeline_cls

Author: jsbueno

extraasync/sync_async_bridge.py

@@ -6,6 +6,8 @@
 import threading
 import typing as t
 
+from asyncio.events import AbstractEventLoop
+from contextvars import ContextVar
 from queue import Queue as ThreadingQueue
 from textwrap import dedent as D
 
@@ -27,19 +29,16 @@
     console_handler.setLevel(logging.DEBUG)
 
     # Create a formatter and attach it to the handler
-    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
     console_handler.setFormatter(formatter)
 
     # Attach the handler to the logger
     logger.addHandler(console_handler)
 
 
 ###########################
-
-T = t.TypeVar("T")
-
-_non_bridge_loop = contextvars.ContextVar("non_bridge_loop", default=None)
-_context_bound_loop = contextvars.ContextVar("_context_bound_loop", default=None)
 _poison = object()
 
 
@@ -52,18 +51,28 @@ class _QueueSet(t.NamedTuple):
 class _SyncTask(t.NamedTuple):
     sync_task: t.Callable
     args: tuple
-    kwargs: dict
-    loop: asyncio.BaseEventLoop
+    kwargs: t.Mapping
+    loop: AbstractEventLoop
     context: contextvars.Context
     done_future: asyncio.Future
 
 
+T = t.TypeVar("T")
+
+_non_bridge_loop: ContextVar[t.Optional[AbstractEventLoop]] = ContextVar(
+    "non_bridge_loop", default=None
+)
+_context_bound_task: ContextVar[t.Optional[_SyncTask]] = ContextVar(
+    "_context_bound_task", default=None
+)
+
+
 def sync_to_async(
-    func: t.Callable[[...,], T] | t.Coroutine,
+    func: t.Callable[..., T] | t.Coroutine,
     args: t.Sequence[t.Any] = (),
-    kwargs: t.Mapping[str, t.Any | None] = None
+    kwargs: t.Optional[t.Mapping[str, t.Any | None]] = None,
 ) -> T:
-    """ Allows calling an async function from a synchronous context.
+    """Allows calling an async function from a synchronous context.
 
     When called from a synchronous scenario, this
     will create and cache an asyncio.loop instance per thread for code that otherwise
@@ -88,9 +97,9 @@ def sync_to_async(
             raise TypeError("Can't accept extra arguments for existing coroutine")
         coro = func
     else:
-        coro = func(*args, **kwargs)
+        coro = t.cast(t.Callable, func)(*args, **kwargs)
 
-    root_sync_task = _context_bound_loop.get()
+    root_sync_task = _context_bound_task.get()
     if not root_sync_task:
         return _sync_to_async_non_bridge(coro, args, kwargs)
 
@@ -100,11 +109,16 @@ def sync_to_async(
     event = _ThreadPool.get(loop).all[threading.current_thread()].event
 
     event.clear()
-    task = None # Keep a strong reference to the task
+    task: None | asyncio.Task = None  # Keep a strong reference to the task
     inner_exception = None
+
     def do_it():
         nonlocal task, inner_exception
-        logger.debug("Creating task in %s from %s", loop, thread_name:=threading.current_thread().name)
+        logger.debug(
+            "Creating task in %s from %s",
+            loop,
+            thread_name := threading.current_thread().name,
+        )
         try:
             task = loop.create_task(coro, context=sync_thread_context_copy)
         except Exception as exc:
@@ -118,37 +132,42 @@ def do_it():
     # Pauses sync-worker thread until original co-routine is finhsed in
     # the original event loop:
     event.wait()
+    if task is None:
+        # should be unreachable - but this check keeps the linters quiet!
+        raise RuntimeError("SyncToAsync error: assynchronous code not executed")
     if inner_exception:
         raise inner_exception
-    if exc:=task.exception():
+    if exc := task.exception():
         raise exc
     return task.result()
 
 
 def _sync_to_async_non_bridge(
-    coro: t.Coroutine,
+    coro: t.Coroutine[t.Any, t.Any, T],
     args: t.Sequence[t.Any],
-    kwargs: t.Mapping[str, t.Any]
+    kwargs: t.Mapping[str, t.Any],
 ) -> T:
     loop = _non_bridge_loop.get()
     if not loop:
         try:
             loop = asyncio.new_event_loop()
         except RuntimeError:
-            raise RuntimeError(D("""\
+            raise RuntimeError(
+                D(
+                    """\
                 Error trying to create a new async loop - to be able to call 'sync_to_async' from
                 code inside a running async loop, the parent 'sync execution branch' must be called
                 with `extraasync.async_to_sync`.
-            """))
+            """
+                )
+            )
         _non_bridge_loop.set(loop)
 
-
     return loop.run_until_complete(coro)
 
 
-
 class _ThreadPool:
-    loop: asyncio.BaseEventLoop
+    loop: AbstractEventLoop
 
     pools = dict()
 
@@ -163,11 +182,12 @@ def get(cls, loop=None):
 
         return cls.pools[loop]
 
-
     def __init__(self):
-        self.idle = set()
+        self.idle = set[_QueueSet | tuple[()]]()
         self.all = dict()
-        self.running = contextvars.ContextVar("running", default=())
+        self.running: ContextVar[_QueueSet | tuple[()]] = ContextVar(
+            "running", default=()
+        )
 
     def __enter__(self):
         if self.idle:
@@ -201,29 +221,31 @@ def __repr__(self):
         return f"<_ThreadPool with {len(self.all)} threads - with {len(self.all) - len(self.idle)} threads in use>"
 
 
-
 def _in_context_sync_worker(sync_task: _SyncTask):
 
-    _context_bound_loop.set(sync_task)
+    _context_bound_task.set(sync_task)
     try:
         logger.debug("Entering sync call in worker thread %s", sync_task)
         result = sync_task.sync_task(*sync_task.args, **sync_task.kwargs)
         logger.debug("Returning from sync call in worker thread: %s", result)
     finally:
-        _context_bound_loop.set(None)
+        _context_bound_task.set(None)
     return result
 
 
 def _sync_worker(queue):
-    """Inner function to call sync "tasks" in a separate thread.
-
-
-    """
+    """Inner function to call sync "tasks" in a separate thread."""
     while True:
         sync_task_bundle = queue.get()
-        logger.debug("*" * 100 + "Got new sync call %s in worker-thread %s", sync_task_bundle, threading.current_thread().name)
+        logger.debug(
+            "*" * 100 + "Got new sync call %s in worker-thread %s",
+            sync_task_bundle,
+            threading.current_thread().name,
+        )
         if sync_task_bundle is _poison:
-            logger.info("Stopping sync-worker thread %s", threading.current_thread().name)
+            logger.info(
+                "Stopping sync-worker thread %s", threading.current_thread().name
+            )
             return
         context = sync_task_bundle.context
         loop = sync_task_bundle.loop
@@ -234,15 +256,13 @@ def _sync_worker(queue):
             result = exc
             loop.call_soon_threadsafe(fut.set_exception, result)
         else:
-            loop.call_soon_threadsafe(fut.set_result, result )
-
+            loop.call_soon_threadsafe(fut.set_result, result)
 
 
 def async_to_sync(
-    func: t.Callable[[...,], T],
-    *,
+    func: t.Callable[..., T],
     args: t.Sequence[t.Any] = (),
-    kwargs: t.Mapping[str, t.Any | None] = None
+    kwargs: t.Optional[t.Mapping[str, t.Any | None]] = None,
 ) -> asyncio.Future[T]:
     """Returns a future wrapping a synchronous call in other thread
 
@@ -269,10 +289,14 @@ def async_to_sync(
         kwargs = {}
 
     task = asyncio.current_task()
+    if task is None:
+        raise RuntimeError("async_to_sync called outside of an asyncio task.")
     loop = task.get_loop()
-    context = task.get_context() # it is resposibility of those using the context to copy it!
-                                 # (also, with the "extracontext" package, there are
-                                 # tools to modify an existing context if needed
+    context = (
+        task.get_context()
+    )  # it is resposibility of those using the context to copy it!
+    # (also, with the "extracontext" package, there are
+    # tools to modify an existing context if needed
 
     done_future = loop.create_future()
 
@@ -286,10 +310,10 @@ def async_to_sync(
         lambda fut, thread_pool=thread_pool: thread_pool.__exit__(None, None, None)
     )
 
-
-    #with _ThreadPool.get() as queue_set:
-    queue_set.queue.put(_SyncTask(func, args, kwargs, loop, context, done_future))
+    # with _ThreadPool.get() as queue_set:
+    queue_set.queue.put(
+        _SyncTask(func, tuple(args), kwargs, loop, context, done_future)
+    )
     logger.debug("Created future awaiting sync result from worker thread for %s", func)
 
     return done_future
-