remove duplicate test + some tweaks

Author: RichardZhu2

src/pyper/_core/sync_helper/output.py

@@ -1,21 +1,22 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-import queue
+from typing import TYPE_CHECKING, Union
 
 from .stage import Producer, ProducerConsumer
 from ..util.sentinel import StopSentinel
 from ..util.worker_pool import ProcessPool, ThreadPool
 
 if TYPE_CHECKING:
+    import multiprocessing as mp
+    import queue
     from ..pipeline import Pipeline
 
 
 class PipelineOutput:
     def __init__(self, pipeline: Pipeline):
         self.pipeline = pipeline
 
-    def _get_q_out(self, tp: ThreadPool, pp: ProcessPool, *args, **kwargs) -> queue.Queue:
+    def _get_q_out(self, tp: ThreadPool, pp: ProcessPool, *args, **kwargs) -> Union[mp.Queue, queue.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]):
@@ -34,5 +35,10 @@ def __call__(self, *args, **kwargs):
         """Iterate through the pipeline, taking the inputs to the first task, and yielding each output from the last task."""
         with ThreadPool() as tp, ProcessPool() as pp:
             q_out = self._get_q_out(tp, pp, *args, **kwargs)
-            while (data := q_out.get()) is not StopSentinel:
-                yield data
+            try:
+                while (data := q_out.get()) is not StopSentinel:
+                    yield data
+            except (KeyboardInterrupt, SystemExit):  # pragma: no cover
+                tp.shutdown_event.set()
+                pp.shutdown_event.set()
+                raise

src/pyper/_core/sync_helper/stage.py

@@ -9,8 +9,8 @@
 from ..util.sentinel import StopSentinel
 
 if TYPE_CHECKING:
+    import multiprocessing as mp
     from multiprocessing.managers import SyncManager
-    import multiprocessing.queues as mpq
     import multiprocessing.synchronize as mpsync
     from ..util.worker_pool import WorkerPool
     from ..task import Task
@@ -53,7 +53,7 @@ def start(self, pool: WorkerPool, /, *args, **kwargs):
 class ProducerConsumer:
     def __init__(
             self,
-            q_in: Union[mpq.Queue, queue.Queue],
+            q_in: Union[mp.Queue, queue.Queue],
             task: Task,
             next_task: Task,
             manager: SyncManager,

tests/test_sync.py

@@ -73,15 +73,6 @@ def test_invalid_branch_result():
         assert isinstance(e, TypeError)
     else:
         raise AssertionError
-    
-def test_invalid_branch():
-    try:
-        p = task(f1, join=True) | task(f2, branch=True) > consumer
-        p(1)
-    except Exception as e:
-        assert isinstance(e, RuntimeError)
-    else:
-        raise AssertionError
 
 def test_threaded_error_handling():
     try: