tidy up misc code

Author: RichardZhu2

src/pyper/_core/async_helper/stage.py

@@ -19,9 +19,9 @@
 class AsyncProducer:
     def __init__(self, task: Task, next_task: Task):
         if task.workers > 1:
-            raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have more than 1 worker")
+            raise RuntimeError(f"The first task in a pipeline ({task.func}) cannot have more than 1 worker")
         if task.join:
-            raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot join previous results")
+            raise RuntimeError(f"The first task in a pipeline ({task.func}) cannot join previous results")
         self.task = task
         self.q_out = asyncio.Queue(maxsize=task.throttle)
 

src/pyper/_core/sync_helper/output.py

@@ -34,11 +34,5 @@ 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 True:
-                try:
-                    # Use the timeout strategy for unblocking main thread without busy waiting
-                    if (data := q_out.get()) is StopSentinel:
-                        break
-                    yield data
-                except (KeyboardInterrupt, SystemExit): # pragma: no cover
-                    raise
+            while (data := q_out.get()) is not StopSentinel:
+                yield data

src/pyper/_core/sync_helper/queue_io.py

@@ -30,13 +30,13 @@ def __call__(self):
 
 
 class _SingleDequeue(_Dequeue):
-     def __call__(self):
-         for data in self._input_stream():
+    def __call__(self):
+        for data in self._input_stream():
             yield data
 
 
 class _JoiningDequeue(_Dequeue):
-     def __call__(self):
+    def __call__(self):
         yield self._input_stream()
 
 
@@ -56,12 +56,12 @@ def __call__(self, *args, **kwargs):
 
 
 class _SingleEnqueue(_Enqueue):        
-     def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs):
         self.q_out.put(self.task.func(*args, **kwargs))
 
 
 class _BranchingEnqueue(_Enqueue):
-     def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs):
         if isinstance(result := self.task.func(*args, **kwargs), Iterable):
             for output in result:
                 self.q_out.put(output)

src/pyper/_core/sync_helper/stage.py

@@ -2,11 +2,11 @@
 
 import queue
 import threading
+from types import SimpleNamespace
 from typing import TYPE_CHECKING, Union
 
 from .queue_io import DequeueFactory, EnqueueFactory
 from ..util.sentinel import StopSentinel
-from ..util.counter import MultiprocessingCounter, ThreadingCounter
 
 if TYPE_CHECKING:
     from multiprocessing.managers import SyncManager
@@ -24,9 +24,9 @@ def __init__(
             manager: SyncManager,
             shutdown_event: Union[mpsync.Event, threading.Event]):
         if task.workers > 1:
-            raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have more than 1 worker")
+            raise RuntimeError(f"The first task in a pipeline ({task.func}) cannot have more than 1 worker")
         if task.join:
-            raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot join previous results")
+            raise RuntimeError(f"The first task in a pipeline ({task.func}) cannot join previous results")
         self.q_out = manager.Queue(maxsize=task.throttle) \
             if task.multiprocess or (next_task is not None and next_task.multiprocess) \
             else queue.Queue(maxsize=task.throttle)
@@ -69,7 +69,8 @@ def __init__(
         self._n_consumers = 1 if next_task is None else next_task.workers
         self._dequeue = DequeueFactory(q_in, task)
         self._enqueue = EnqueueFactory(self.q_out, task)
-        self._workers_done = MultiprocessingCounter(0, manager) if task.multiprocess else ThreadingCounter(0)
+        self._workers_done = manager.Value('i', 0) if task.multiprocess else SimpleNamespace(value=0)
+        self._workers_done_lock = manager.Lock() if task.multiprocess else threading.Lock()
 
     def _worker(self):
         try:
@@ -80,11 +81,12 @@ def _worker(self):
             self._shutdown_event.set()
             raise
         finally:
-            if self._workers_done.increment() == self._n_workers:
-                for _ in range(self._n_consumers):
-                    self.q_out.put(StopSentinel)
+            with self._workers_done_lock:
+                self._workers_done.value += 1
+                if self._workers_done.value == self._n_workers:
+                    for _ in range(self._n_consumers):
+                        self.q_out.put(StopSentinel)
 
-            
     def start(self, pool: WorkerPool, /):
         for _ in range(self._n_workers):
             pool.submit(self._worker)

src/pyper/_core/util/counter.py

@@ -27,9 +27,6 @@ def __exit__(self, et, ev, tb):
         for future in self._futures:
             future.result()
 
-    def create_queue(self):
-        raise NotImplementedError
-
     def submit(self, func, /, *args, **kwargs):
         future = self._executor.submit(func, *args, **kwargs)
         self._futures.append(future)