implement explicit branch parameter (BREAKING CHANGE)

Author: RichardZhu2

src/pyper/_core/async_helper/queue_io.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import AsyncIterable, Iterable
 from typing import TYPE_CHECKING
 
 from ..util.sentinel import StopSentinel
@@ -39,7 +40,7 @@ async def __call__(self):
 
 
 def AsyncEnqueueFactory(q_out: asyncio.Queue, task: Task):
-    return _BranchingAsyncEnqueue(q_out=q_out, task=task) if task.is_gen \
+    return _BranchingAsyncEnqueue(q_out=q_out, task=task) if task.branch \
         else _SingleAsyncEnqueue(q_out=q_out, task=task)
 
 
@@ -60,5 +61,13 @@ async def __call__(self, *args, **kwargs):
 
 class _BranchingAsyncEnqueue(_AsyncEnqueue):
     async def __call__(self, *args, **kwargs):
-        async for output in self.task.func(*args, **kwargs):
-            await self.q_out.put(output)
+        result = self.task.func(*args, **kwargs)
+        if isinstance(result, AsyncIterable):
+            async for output in result:
+                await self.q_out.put(output)
+        elif isinstance(data := await result, Iterable):
+            for output in data:
+                await self.q_out.put(output)
+        else:
+            raise TypeError(f"got object of type {type(data)} 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/decorators.py

@@ -19,10 +19,11 @@
 
 
 class task:
-    """Decorator class to initialize a `Pipeline` consisting of one task, from a function or callable.
+    """Decorator class to initialize a `Pipeline` consisting of one task.
 
     Args:
-        func (callable): A positional-only param defining the task function
+        func (callable): A positional-only param defining the task function (can be omitted when using `@task`)
+        branch (bool): Allows the task to submit multiple outputs
         join (bool): Allows the task to take all previous results as input, instead of single results
         workers (int): Defines the number of workers to run the task
         throttle (int): Limits the number of results the task is able to produce when all consumers are busy
@@ -32,7 +33,7 @@ class task:
     Returns:
         Pipeline: A `Pipeline` instance consisting of one task.
 
-     Example:
+    Example:
     ```python
     def f(x: int):
         return x + 1
@@ -46,6 +47,7 @@ def __new__(
             func: None = None,
             /,
             *,
+            branch: bool = False,
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
@@ -55,21 +57,23 @@ def __new__(
     @t.overload
     def __new__(
             cls,
-            func: t.Callable[_P, t.Awaitable[_R]],
+            func: t.Callable[_P, t.Union[t.Awaitable[t.Iterable[_R]], t.AsyncGenerator[_R]]],
             /,
             *,
+            branch: True,
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
             bind: _ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
-    
+        
     @t.overload
     def __new__(
             cls,
-            func: t.Callable[_P, t.AsyncGenerator[_R]],
+            func: t.Callable[_P, t.Awaitable[_R]],
             /,
             *,
+            branch: bool = False,
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
@@ -79,9 +83,10 @@ def __new__(
     @t.overload
     def __new__(
             cls,
-            func: t.Callable[_P, t.Generator[_R]],
+            func: t.Callable[_P, t.Iterable[_R]],
             /,
             *,
+            branch: True,
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
@@ -94,6 +99,7 @@ def __new__(
             func: t.Callable[_P, _R],
             /,
             *,
+            branch: bool = False,
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
@@ -105,15 +111,16 @@ def __new__(
             func: t.Optional[t.Callable] = None,
             /,
             *,
+            branch: bool = False,
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
             bind: _ArgsKwargs = None):
         # Classic decorator trick: @task() means func is None, @task without parentheses means func is passed. 
         if func is None:
-            return functools.partial(cls, join=join, workers=workers, throttle=throttle, multiprocess=multiprocess, bind=bind)
-        return Pipeline([Task(func=func, join=join, workers=workers, throttle=throttle, multiprocess=multiprocess, bind=bind)])
+            return functools.partial(cls, branch=branch, join=join, workers=workers, throttle=throttle, multiprocess=multiprocess, bind=bind)
+        return Pipeline([Task(func=func, branch=branch, join=join, workers=workers, throttle=throttle, multiprocess=multiprocess, bind=bind)])
 
     @staticmethod
     def bind(*args, **kwargs) -> _ArgsKwargs:
@@ -131,4 +138,3 @@ def f(x: int, y: int):
         if not args and not kwargs:
             return None
         return args, kwargs
-    

src/pyper/_core/sync_helper/queue_io.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Iterable
 from typing import TYPE_CHECKING, Union
 
 from ..util.sentinel import StopSentinel
@@ -40,7 +41,7 @@ def __call__(self):
 
 
 def EnqueueFactory(q_out: Union[mp.Queue, queue.Queue], task: Task):
-    return _BranchingEnqueue(q_out=q_out, task=task) if task.is_gen \
+    return _BranchingEnqueue(q_out=q_out, task=task) if task.branch \
         else _SingleEnqueue(q_out=q_out, task=task)
 
 
@@ -61,5 +62,11 @@ def __call__(self, *args, **kwargs):
 
 class _BranchingEnqueue(_Enqueue):
      def __call__(self, *args, **kwargs):
-        for output in self.task.func(*args, **kwargs):
-            self.q_out.put(output)
+        result = self.task.func(*args, **kwargs)
+        if isinstance(result, Iterable):
+            for output in result:
+                self.q_out.put(output)
+        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/task.py

@@ -9,18 +9,20 @@ class Task:
     """The representation of a function within a Pipeline."""
 
     __slots__ = (
-        "is_gen",
-        "is_async",
         "func",
+        "branch",
         "join",
         "workers",
         "throttle",
-        "multiprocess"
+        "multiprocess",
+        "is_async",
+        "is_gen"
     )
 
     def __init__(
             self,
             func: Callable,
+            branch: bool = False,
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
@@ -35,7 +37,7 @@ def __init__(
         if throttle < 0:
             raise ValueError("throttle cannot be less than 0")
         if not callable(func):
-            raise TypeError("A task must be a callable object")
+            raise TypeError("A task function must be a callable object")
 
         self.is_gen = inspect.isgeneratorfunction(func) \
             or inspect.isasyncgenfunction(func) \
@@ -50,6 +52,7 @@ def __init__(
             raise ValueError("multiprocess cannot be True for an async task")
 
         self.func = func if bind is None else functools.partial(func, *bind[0], **bind[1])
+        self.branch = branch
         self.join = join
         self.workers = workers
         self.throttle = throttle

src/pyper/_core/util/asynchronize.py

@@ -17,7 +17,7 @@ def ascynchronize(task: Task, tp: ThreadPoolExecutor, pp: ProcessPoolExecutor) -
     if task.is_async:
         return task
 
-    if task.is_gen:
+    if task.is_gen and task.branch:
         @functools.wraps(task.func)
         async def wrapper(*args, **kwargs):
             for output in task.func(*args, **kwargs):
@@ -31,6 +31,7 @@ async def wrapper(*args, **kwargs):
             return await loop.run_in_executor(executor=executor, func=f)
     return Task(
         func=wrapper,
+        branch=task.branch,
         join=task.join,
         workers=task.workers,
         throttle=task.throttle