update tests, along with small refactors

Author: RichardZhu2

src/pyper/_core/async_helper/queue_io.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterable, Iterable
+from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
 from ..util.sentinel import StopSentinel
@@ -61,13 +61,12 @@ async def __call__(self, *args, **kwargs):
 
 class _BranchingAsyncEnqueue(_AsyncEnqueue):
     async def __call__(self, *args, **kwargs):
-        result = self.task.func(*args, **kwargs)
-        if isinstance(result, AsyncIterable):
-            async for output in result:
+        if self.task.is_gen:
+            async for output in self.task.func(*args, **kwargs):
                 await self.q_out.put(output)
-        elif isinstance(data := await result, Iterable):
-            for output in data:
+        elif isinstance(result := await self.task.func(*args, **kwargs), Iterable):
+            for output in result:
                 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"
+            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/sync_helper/output.py

@@ -45,7 +45,7 @@ def __call__(self, *args, **kwargs):
                 except queue.Empty:
                     tp.raise_error_if_exists()
                     pp.raise_error_if_exists()
-                except (KeyboardInterrupt, SystemExit):
+                except (KeyboardInterrupt, SystemExit): # pragma: no cover
                     tp.shutdown_event.set()
                     pp.shutdown_event.set()
                     raise

src/pyper/_core/sync_helper/queue_io.py

@@ -62,8 +62,7 @@ def __call__(self, *args, **kwargs):
 
 class _BranchingEnqueue(_Enqueue):
      def __call__(self, *args, **kwargs):
-        result = self.task.func(*args, **kwargs)
-        if isinstance(result, Iterable):
+        if isinstance(result := self.task.func(*args, **kwargs), Iterable):
             for output in result:
                 self.q_out.put(output)
         else:

src/pyper/_core/sync_helper/stage.py

@@ -2,14 +2,15 @@
 
 import multiprocessing as mp
 import queue
-import threading
 from typing import TYPE_CHECKING, Union
 
 from .queue_io import DequeueFactory, EnqueueFactory
 from ..util.sentinel import StopSentinel
+from ..util.value import ThreadingValue
 
 if TYPE_CHECKING:
-    from multiprocessing.synchronize import Event as MpEvent
+    from multiprocessing.synchronize import Event as ProcessEvent
+    from threading import Event as ThreadEvent
     from ..util.worker_pool import WorkerPool
     from ..task import Task
 
@@ -20,7 +21,7 @@ def __init__(
             task: Task,
             next_task: Task,
             q_err: Union[mp.Queue, queue.Queue],
-            shutdown_event: Union[MpEvent, threading.Event]):
+            shutdown_event: Union[ProcessEvent, ThreadEvent]):
         if task.workers > 1:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have more than 1 worker")
         if task.join:
@@ -56,7 +57,7 @@ def __init__(
             task: Task,
             next_task: Task,
             q_err: Union[mp.Queue, queue.Queue],
-            shutdown_event: Union[MpEvent, threading.Event]):
+            shutdown_event: Union[ProcessEvent, ThreadEvent]):
         # The output queue is shared between this task and the next. We optimize here by using queue.Queue wherever possible
         # and only using multiprocess.Queue when the current task or the next task are multiprocessed
         self.q_out = mp.Queue(maxsize=task.throttle) \
@@ -69,23 +70,12 @@ 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._multiprocess = task.multiprocess
-        if self._multiprocess:
-            self._workers_done = mp.Value('i', 0)
-            self._lock = self._workers_done.get_lock()
-        else:
-            self._workers_done = 0
-            self._lock = threading.Lock()
+        self._workers_done = mp.Value('i', 0) if task.multiprocess else ThreadingValue(0)
 
     def _increment_workers_done(self):
-        with self._lock:
-            if self._multiprocess:
-                self._workers_done.value += 1
-                return self._workers_done.value
-            else:
-                self._workers_done += 1
-                return self._workers_done
+        with self._workers_done.get_lock():
+            self._workers_done.value += 1
+            return self._workers_done.value
 
     def _finish(self):
         if self._increment_workers_done() == self._n_workers:

src/pyper/_core/util/value.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+
+class ThreadingValue:
+    """Utility class to help manage thread based access to a value.
+    
+    The `get_lock` method mimics the API of `multiprocessing.Value`
+    """
+    def __init__(self, value: Any):
+        self.value = value
+        self._lock = threading.Lock()
+    
+    def get_lock(self):
+        return self._lock

src/pyper/_core/util/worker_pool.py

@@ -16,7 +16,7 @@ class WorkerPool:
     worker_type = None
 
     def __init__(self):
-        self.error_queue = mp.Queue(1) if self.worker_type is mp.Process else queue.Queue(1)
+        self.error_queue = mp.Queue() if self.worker_type is mp.Process else queue.Queue()
         self.shutdown_event = mp.Event() if self.worker_type is mp.Process else threading.Event()
 
         self._workers: List[Union[mp.Process, threading.Thread]] = []
@@ -35,9 +35,6 @@ def has_error(self):
     def get_error(self) -> Exception:
         return self.error_queue.get()
 
-    def put_error(self, e: Exception):
-        self.error_queue.put(e)
-
     def raise_error_if_exists(self):
         if self.has_error:
             raise self.get_error() from None

tests/Dockerfile

@@ -13,4 +13,4 @@ COPY .git .git
 
 RUN chmod +x tests/entrypoint.sh
 
-ENTRYPOINT ["tests/entrypoint.sh"]
+ENTRYPOINT ["/bin/bash", "tests/entrypoint.sh"]

tests/test_async.py

@@ -1,6 +1,7 @@
 from pyper import task
 import pytest
 
+class TestError(Exception): ...
 
 def f1(data):
     return data
@@ -9,7 +10,10 @@ def f2(data):
     yield data
 
 def f3(data):
-    raise RuntimeError
+    raise TestError
+
+def f4(data):
+    return [data]
 
 async def af1(data):
     return data
@@ -18,7 +22,7 @@ async def af2(data):
     yield data
 
 async def af3(data):
-    raise RuntimeError
+    raise TestError
 
 async def af4(data):
     async for row in data:
@@ -31,24 +35,29 @@ async def consumer(data):
     return total
 
 @pytest.mark.asyncio
-async def test_pipeline():
-    p = task(f1) | task(f2)
-    assert p(1).__next__() == 1
+async def test_aiterable_branched_pipeline():
+    p = task(af1) | task(f2, branch=True)
+    assert await p(1).__anext__() == 1
+
+@pytest.mark.asyncio
+async def test_iterable_branched_pipeline():
+    p = task(af1) | task(f4, branch=True)
+    assert await p(1).__anext__() == 1
 
 @pytest.mark.asyncio
 async def test_joined_pipeline():
-    p = task(af1) | task(af2) | task(af4, join=True)
+    p = task(af1) | task(af2, branch=True) | task(af4, branch=True, join=True)
     assert await p(1).__anext__() == 1
 
 @pytest.mark.asyncio
 async def test_consumer():
-    p = task(af1) | task(af2) > consumer
+    p = task(af1) | task(af2, branch=True) > consumer
     assert await p(1) == 1
 
 @pytest.mark.asyncio
-async def test_invalid_first_stage_concurrency():
+async def test_invalid_first_stage_workers():
     try:
-        p = task(af1, concurrency=2) | task(af2) > consumer
+        p = task(af1, workers=2) | task(af2, branch=True) > consumer
         await p(1)
     except Exception as e:
         assert isinstance(e, RuntimeError)
@@ -58,35 +67,48 @@ async def test_invalid_first_stage_concurrency():
 @pytest.mark.asyncio
 async def test_invalid_first_stage_join():
     try:
-        p = task(af1, join=True) | task(af2) > consumer
+        p = task(af1, join=True) | task(af2, branch=True) > consumer
         await p(1)
     except Exception as e:
         assert isinstance(e, RuntimeError)
     else:
         raise AssertionError
-
+    
 @pytest.mark.asyncio
-async def test_error_handling():
+async def test_invalid_branch_result():
     try:
-        p = task(af1) | task(af2) | task(af3) > consumer
+        p = task(af1, branch=True) > consumer
         await p(1)
     except Exception as e:
-        assert isinstance(e, RuntimeError)
+        assert isinstance(e, TypeError)
     else:
         raise AssertionError
-    
-@pytest.mark.asyncio
-async def test_unified_pipeline():
-    p = task(af1) | task(f1) | task(af2) | task(f2) > consumer
-    assert await p(1) == 1
 
-@pytest.mark.asyncio
-async def test_error_handling_in_daemon():
+async def _try_catch_error(pipeline):
     try:
-        p = task(af1) | task(af2) | task(f3, daemon=True) > consumer
+        p = task(af1) | pipeline > consumer
         await p(1)
     except Exception as e:
-        assert isinstance(e, RuntimeError)
+        return isinstance(e, TestError)
     else:
-        raise AssertionError
-    
+        return False
+    
+@pytest.mark.asyncio
+async def test_async_error_handling():
+    p = task(af3)
+    assert await _try_catch_error(p)
+    
+@pytest.mark.asyncio
+async def test_threaded_error_handling():
+    p = task(f3, workers=2)
+    assert await _try_catch_error(p)
+    
+@pytest.mark.asyncio
+async def test_multiprocessed_error_handling():
+    p = task(f3, workers=2, multiprocess=True)
+    assert await _try_catch_error(p)
+    
+@pytest.mark.asyncio
+async def test_unified_pipeline():
+    p = task(af1) | task(f1) | task(f2, branch=True, multiprocess=True) > consumer
+    assert await p(1) == 1

tests/test_sync.py

@@ -1,6 +1,7 @@
 import time
 from pyper import task
 
+class TestError(Exception): ...
 
 def f1(data):
     return data
@@ -18,37 +19,37 @@ def f4(a1, a2, a3, data, k1, k2):
 def f5(data):
     # Make queue monitor timeout on main thread
     time.sleep(0.2)
-    raise RuntimeError
+    raise TestError
 
 def consumer(data):
     total = 0
     for i in data:
         total += i
     return total
 
-def test_pipeline():
-    p = task(f1) | task(f2)
+def test_branched_pipeline():
+    p = task(f1) | task(f2, branch=True)
     assert p(1).__next__() == 1
 
 def test_joined_pipeline():
-    p = task(f1) | task(f2) | task(f3, join=True)
+    p = task(f1) | task(f2, branch=True) | task(f3, branch=True, join=True)
     assert p(1).__next__() == 1
 
 def test_bind():
     p = task(f1) | task(f4, bind=task.bind(1, 1, 1, k1=1, k2=2))
     assert p(1).__next__() == 1
 
 def test_redundant_bind_ok():
-    p = task(f1) | task(f2, bind=task.bind())
+    p = task(f1) | task(f2, branch=True, bind=task.bind())
     assert p(1).__next__() == 1
 
 def test_consumer():
-    p = task(f1) | task(f2) > consumer
+    p = task(f1) | task(f2, branch=True) > consumer
     assert p(1) == 1
 
-def test_invalid_first_stage_concurrency():
+def test_invalid_first_stage_workers():
     try:
-        p = task(f1, concurrency=2) | task(f2) > consumer
+        p = task(f1, workers=2) | task(f2) > consumer
         p(1)
     except Exception as e:
         assert isinstance(e, RuntimeError)
@@ -57,29 +58,45 @@ def test_invalid_first_stage_concurrency():
 
 def test_invalid_first_stage_join():
     try:
-        p = task(f1, join=True) | task(f2) > consumer
+        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_error_handling():
+def test_invalid_branch_result():
     try:
-        p = task(f1) | task(f2) | task(f5) > consumer
+        p = task(f1, branch=True) > consumer
+        p(1)
+    except Exception as e:
+        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:
-        print(e)
         assert isinstance(e, RuntimeError)
     else:
         raise AssertionError
 
-def test_error_handling_in_daemon():
+def test_threaded_error_handling():
     try:
-        p = task(f5, daemon=True) | task(f2) > consumer
+        p = task(f1) | task(f5, workers=2) > consumer
         p(1)
     except Exception as e:
-        print(e)
-        assert isinstance(e, RuntimeError)
+        assert isinstance(e, TestError)
+    else:
+        raise AssertionError
+
+def test_multiprocessed_error_handling():
+    try:
+        p = task(f1) | task(f5, workers=2, multiprocess=True) > consumer
+        p(1)
+    except Exception as e:
+        assert isinstance(e, TestError)
     else:
         raise AssertionError