move unblocking logic to queue_io

Author: RichardZhu2

src/pyper/_core/async_helper/queue_io.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from collections.abc import AsyncIterable, Iterable
 from typing import TYPE_CHECKING
 
@@ -65,9 +66,11 @@ async def __call__(self, *args, **kwargs):
         if isinstance(result, AsyncIterable):
             async for output in result:
                 await self.q_out.put(output)
+                await asyncio.sleep(0)
         elif isinstance(result := await result, Iterable):
             for output in result:
                 await self.q_out.put(output)
+                await asyncio.sleep(0)
         else:
             raise TypeError(f"got object of type {type(result)} from branching task {self.task.func} which could not be iterated over"
                             " (the task should be a generator, or return an iterable)")

src/pyper/_core/util/asynchronize.py

@@ -10,22 +10,18 @@
 def ascynchronize(task: Task, tp: ThreadPoolExecutor, pp: ProcessPoolExecutor) -> Task:
     """Unify async and sync tasks as awaitable futures.
     1. If the task is async already, return it.
-    2. Synchronous generators are transformed into asynchronous generators.
-    3. Multiprocessed synchronous functions are wrapped in a call to `run_in_executor` using `ProcessPoolExecutor`.
-    4. Threaded synchronous functions are wrapped in a call to `run_in_executor` using `ThreadPoolExecutor`.
+    2. Multiprocessed synchronous functions are wrapped in a call to `run_in_executor` using `ProcessPoolExecutor`.
+    3. Threaded synchronous functions are wrapped in a call to `run_in_executor` using `ThreadPoolExecutor`.
     """
     if task.is_async:
         return task
 
-    if task.is_gen and task.branch:
-        # Small optimization to convert sync generators to async generators
+    if task.is_gen:
+        # Small optimization to convert sync generators to async functions
         # This saves from having to use a thread/process just to get the generator object
-        # We also add asyncio.sleep(0) to unblock long synchronous generators
         @functools.wraps(task.func)
         async def wrapper(*args, **kwargs):
-            for output in task.func(*args, **kwargs):
-                yield output
-                await asyncio.sleep(0)
+            return task.func(*args, **kwargs)
     else:
         executor = pp if task.multiprocess else tp
         @functools.wraps(task.func)