remove daemon support and refactor async handlers

Author: RichardZhu2

src/pyper/_core/async_helper/output.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import asyncio
-from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 import sys
 from typing import TYPE_CHECKING
 
 from .stage import AsyncProducer, AsyncProducerConsumer
+from ..util.asynchronize import ascynchronize
 from ..util.sentinel import StopSentinel
-from ..util.thread_pool import ThreadPool
 
 if sys.version_info < (3, 11):  # pragma: no cover
     from ..util.task_group import TaskGroup, ExceptionGroup
@@ -22,30 +22,32 @@ class AsyncPipelineOutput:
     def __init__(self, pipeline: AsyncPipeline):
         self.pipeline = pipeline
 
-    def _get_q_out(self, tg: TaskGroup, tp: ThreadPool, pp: ProcessPoolExecutor, *args, **kwargs) -> asyncio.Queue:
+    def _get_q_out(self, tg: TaskGroup, tp: ThreadPoolExecutor, pp: ProcessPoolExecutor, *args, **kwargs) -> asyncio.Queue:
         """Feed forward each stage to the next, returning the output queue of the final stage."""
         q_out = None
         for task, next_task in zip(self.pipeline.tasks, self.pipeline.tasks[1:] + [None]):
-            n_consumers = 1 if next_task is None else next_task.concurrency
+            task = ascynchronize(task, tp=tp, pp=pp)
             if q_out is None:
-                stage = AsyncProducer(task=self.pipeline.tasks[0], tg=tg, tp=tp, pp=pp, n_consumers=n_consumers)
-                stage.start(*args, **kwargs)
+                stage = AsyncProducer(task=task, next_task=next_task)
+                stage.start(tg, *args, **kwargs)
             else:
-                stage = AsyncProducerConsumer(q_in=q_out, task=task, tg=tg, tp=tp, pp=pp, n_consumers=n_consumers)
-                stage.start()
+                stage = AsyncProducerConsumer(q_in=q_out, task=task, next_task=next_task)
+                stage.start(tg)
             q_out = stage.q_out
 
         return q_out
 
     async def __call__(self, *args, **kwargs):
-        """Call the pipeline, taking the inputs to the first task, and returning the output from the last task."""
+        """Iterate through the pipeline, taking the inputs to the first task, and yielding each output from the last task.
+
+        Unify async, threaded, and multiprocessed work by:
+        1. using TaskGroup to execute asynchronous tasks
+        2. using ThreadPoolExecutor to execute threaded synchronous tasks
+        3. using ProcessPoolExecutor to execute multiprocessed synchronous tasks
+        """
         try:
-            # Unify async, threaded, and multiprocessed work by:
-            # 1. using TaskGroup to execute asynchronous tasks
-            # 2. using ThreadPool to execute threaded synchronous tasks
-            # 3. using ProcessPoolExecutor to execute multiprocessed synchronous tasks
             async with TaskGroup() as tg:
-                with ThreadPool() as tp, ProcessPoolExecutor() as pp:
+                with ThreadPoolExecutor() as tp, ProcessPoolExecutor() as pp:
                     q_out = self._get_q_out(tg, tp, pp, *args, **kwargs)
                     while (data := await q_out.get()) is not StopSentinel:
                         yield data

src/pyper/_core/async_helper/queue_io.py

@@ -9,19 +9,15 @@
     from ..task import Task
 
 
-class AsyncDequeue:
+def AsyncDequeueFactory(q_in: asyncio.Queue, task: Task):
+    return _JoiningAsyncDequeue(q_in=q_in) if task.join \
+        else _SingleAsyncDequeue(q_in=q_in)
+
+
+class _AsyncDequeue:
     """Pulls data from an input queue."""
-    def __new__(self, q_in: asyncio.Queue, task: Task):
-        if task.join:
-            instance = object.__new__(_JoiningAsyncDequeue)
-        else:
-            instance = object.__new__(_SingleAsyncDequeue)
-        instance.__init__(q_in=q_in, task=task)
-        return instance
-
-    def __init__(self, q_in: asyncio.Queue, task: Task):
-         self.q_in = q_in
-         self.task = task
+    def __init__(self, q_in: asyncio.Queue):
+        self.q_in = q_in
 
     async def _input_stream(self):
         while (data := await self.q_in.get()) is not StopSentinel:
@@ -31,41 +27,38 @@ def __call__(self):
         raise NotImplementedError
 
 
-class _SingleAsyncDequeue(AsyncDequeue):
+class _SingleAsyncDequeue(_AsyncDequeue):
     async def __call__(self):
         async for data in self._input_stream():
             yield data
 
 
-class _JoiningAsyncDequeue(AsyncDequeue):
+class _JoiningAsyncDequeue(_AsyncDequeue):
     async def __call__(self):
         yield self._input_stream()
 
 
-class AsyncEnqueue:
-    """Puts output from a task onto an output queue."""
-    def __new__(cls, q_out: asyncio.Queue, task: Task):
-        if task.is_gen:
-            instance = object.__new__(_BranchingAsyncEnqueue)
-        else:
-            instance = object.__new__(_SingleAsyncEnqueue)
-        instance.__init__(q_out=q_out, task=task)
-        return instance
+def AsyncEnqueueFactory(q_out: asyncio.Queue, task: Task):
+    return _BranchingAsyncEnqueue(q_out=q_out, task=task) if task.is_gen \
+        else _SingleAsyncEnqueue(q_out=q_out, task=task)
 
+
+class _AsyncEnqueue:
+    """Puts output from a task onto an output queue."""
     def __init__(self, q_out: asyncio.Queue, task: Task):
-         self.q_out = q_out
-         self.task = task
+        self.q_out = q_out
+        self.task = task
 
     async def __call__(self, *args, **kwargs):
         raise NotImplementedError
 
 
-class _SingleAsyncEnqueue(AsyncEnqueue):        
+class _SingleAsyncEnqueue(_AsyncEnqueue):        
     async def __call__(self, *args, **kwargs):
         await self.q_out.put(await self.task.func(*args, **kwargs))
 
 
-class _BranchingAsyncEnqueue(AsyncEnqueue):
+class _BranchingAsyncEnqueue(_AsyncEnqueue):
     async def __call__(self, *args, **kwargs):
         async for output in self.task.func(*args, **kwargs):
             await self.q_out.put(output)

src/pyper/_core/async_helper/stage.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 
 import asyncio
-from concurrent.futures import ProcessPoolExecutor
 import sys
 from typing import TYPE_CHECKING
 
-from .queue_io import AsyncDequeue, AsyncEnqueue
-from ..util.asynchronize import ascynchronize
+from .queue_io import AsyncDequeueFactory, AsyncEnqueueFactory
 from ..util.sentinel import StopSentinel
 
 if sys.version_info < (3, 11):  # pragma: no cover
@@ -15,67 +13,50 @@
     from asyncio import TaskGroup
 
 if TYPE_CHECKING:
-    from ..util.thread_pool import ThreadPool
     from ..task import Task
 
 
 class AsyncProducer:
-    def __init__(
-            self,
-            task: Task,
-            tg: TaskGroup,
-            tp: ThreadPool,
-            pp: ProcessPoolExecutor,
-            n_consumers: int):
-        self.task = ascynchronize(task, tp, pp)
+    def __init__(self, task: Task, next_task: Task):
         if task.concurrency > 1:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have concurrency greater than 1")
         if task.join:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot join previous results")
-        self.tg = tg
-        self.n_consumers = n_consumers
+        self.task = task
         self.q_out = asyncio.Queue(maxsize=task.throttle)
 
-        self._enqueue = AsyncEnqueue(self.q_out, self.task)
+        self._n_consumers = 1 if next_task is None else next_task.concurrency
+        self._enqueue = AsyncEnqueueFactory(self.q_out, self.task)
 
     async def _worker(self, *args, **kwargs):
         await self._enqueue(*args, **kwargs)
 
-        for _ in range(self.n_consumers):
+        for _ in range(self._n_consumers):
             await self.q_out.put(StopSentinel)
 
-    def start(self, *args, **kwargs):
-        self.tg.create_task(self._worker(*args, **kwargs))
+    def start(self, tg: TaskGroup, /, *args, **kwargs):
+        tg.create_task(self._worker(*args, **kwargs))
 
 
 class AsyncProducerConsumer:
-    def __init__(
-            self,
-            q_in: asyncio.Queue,
-            task: Task,
-            tg: TaskGroup,
-            tp: ThreadPool,
-            pp: ProcessPoolExecutor,
-            n_consumers: int):
-        self.q_in = q_in
-        self.task = ascynchronize(task, tp, pp)
-        self.tg = tg
-        self.n_consumers = n_consumers
+    def __init__(self, q_in: asyncio.Queue, task: Task, next_task: Task):
         self.q_out = asyncio.Queue(maxsize=task.throttle)
 
+        self._n_workers = task.concurrency
+        self._n_consumers = 1 if next_task is None else next_task.concurrency
+        self._dequeue = AsyncDequeueFactory(q_in, task)
+        self._enqueue = AsyncEnqueueFactory(self.q_out, task)
         self._workers_done = 0
-        self._dequeue = AsyncDequeue(self.q_in, self.task)
-        self._enqueue = AsyncEnqueue(self.q_out, self.task)
 
     async def _worker(self):
         async for output in self._dequeue():
             await self._enqueue(output)
 
         self._workers_done += 1
-        if self._workers_done == self.task.concurrency:
-            for _ in range(self.n_consumers):
+        if self._workers_done == self._n_workers:
+            for _ in range(self._n_consumers):
                 await self.q_out.put(StopSentinel)
 
-    def start(self):
-        for _ in range(self.task.concurrency):
-            self.tg.create_task(self._worker())
+    def start(self, tg: TaskGroup, /):
+        for _ in range(self._n_workers):
+            tg.create_task(self._worker())

src/pyper/_core/decorators.py

@@ -19,14 +19,12 @@
 
 
 class task:
-    """Decorator class to transform a function into a `Task` object, and then initialize a `Pipeline` with this task.
-    A Pipeline initialized in this way consists of one Task, and can be piped into other Pipelines.
+    """Decorator class to initialize a `Pipeline` consisting of one task.
 
     The behaviour of each task within a Pipeline is determined by the parameters:
     * `join`: allows the function to take all previous results as input, instead of single results
     * `concurrency`: runs the functions with multiple (async or threaded) workers
     * `throttle`: limits the number of results the function is able to produce when all consumers are busy
-    * `daemon`: determines whether threaded workers are daemon threads (cannot be True for async tasks)
     * `bind`: additional args and kwargs to bind to the function when defining a pipeline
     """
     @t.overload
@@ -38,7 +36,6 @@ def __new__(
             join: bool = False,
             concurrency: int = 1,
             throttle: int = 0,
-            daemon: bool = False,
             multiprocess: bool = False,
             bind: _ArgsKwargs = None) -> t.Type[task]: ...
 
@@ -51,7 +48,6 @@ def __new__(
             join: bool = False,
             concurrency: int = 1,
             throttle: int = 0,
-            daemon: bool = False,
             multiprocess: bool = False,
             bind: _ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
 
@@ -64,7 +60,6 @@ def __new__(
             join: bool = False,
             concurrency: int = 1,
             throttle: int = 0,
-            daemon: bool = False,
             multiprocess: bool = False,
             bind: _ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
 
@@ -77,7 +72,6 @@ def __new__(
             join: bool = False,
             concurrency: int = 1,
             throttle: int = 0,
-            daemon: bool = False,
             multiprocess: bool = False,
             bind: _ArgsKwargs = None) -> Pipeline[_P, _R]: ...
 
@@ -90,7 +84,6 @@ def __new__(
             join: bool = False,
             concurrency: int = 1,
             throttle: int = 0,
-            daemon: bool = False,
             multiprocess: bool = False,
             bind: _ArgsKwargs = None) -> Pipeline[_P, _R]: ...
 
@@ -102,17 +95,25 @@ def __new__(
             join: bool = False,
             concurrency: int = 1,
             throttle: int = 0,
-            daemon: bool = False,
             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, concurrency=concurrency, throttle=throttle, daemon=daemon, multiprocess=multiprocess, bind=bind)
-        return Pipeline([Task(func=func, join=join, concurrency=concurrency, throttle=throttle, daemon=daemon, multiprocess=multiprocess, bind=bind)])
+            return functools.partial(cls, join=join, concurrency=concurrency, throttle=throttle, multiprocess=multiprocess, bind=bind)
+        return Pipeline([Task(func=func, join=join, concurrency=concurrency, throttle=throttle, multiprocess=multiprocess, bind=bind)])
 
     @staticmethod
     def bind(*args, **kwargs) -> _ArgsKwargs:
-        """Utility method, to be used with `functools.partial`."""
+        """Bind additional `args` and `kwargs` to a task.
+
+        Example:
+        ```python
+        def f(x: int, y: int):
+            return x + y
+
+        p = task(f, bind=task.bind(y=1))
+        p(x=1)
+        """
         if not args and not kwargs:
             return None
         return args, kwargs

src/pyper/_core/task.py

@@ -15,7 +15,6 @@ class Task:
         "join",
         "concurrency",
         "throttle",
-        "daemon",
         "multiprocess"
     )
 
@@ -25,7 +24,6 @@ def __init__(
             join: bool = False,
             concurrency: int = 1,
             throttle: int = 0,
-            daemon: bool = False,
             multiprocess: bool = False,
             bind: Optional[Tuple[Tuple, Dict]] = None):
         if not isinstance(concurrency, int):
@@ -38,8 +36,6 @@ def __init__(
             raise ValueError("throttle cannot be less than 0")
         if not callable(func):
             raise TypeError("A task must be a callable object")
-        if daemon and multiprocess:
-            raise ValueError("daemon and multiprocess cannot both be True")
 
         self.is_gen = inspect.isgeneratorfunction(func) \
             or inspect.isasyncgenfunction(func) \
@@ -50,14 +46,11 @@ def __init__(
             or inspect.iscoroutinefunction(func.__call__) \
             or inspect.isasyncgenfunction(func.__call__)
 
-        if self.is_async and daemon:
-            raise ValueError("daemon cannot be True for an async task")
         if self.is_async and multiprocess:
             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.join = join
         self.concurrency = concurrency
         self.throttle = throttle
-        self.daemon = daemon
         self.multiprocess = multiprocess