fix(util): honour flush_at_end and progress_bar in parallelize

- route each worker thread's stdout to its own buffer instead of rebinding sys.stdout
- flush_at_end replays the collected output in input order once the batch ends
- progress_bar=False now selects asyncio.gather over tqdm.gather
- drop the unreachable dask branch that still held the only uses of both
- cover parallelize with 22 offline tests, closes #25 and #154

Author: LukasGold

src/osw/utils/util.py

@@ -1,6 +1,7 @@
 import asyncio
 import functools
 import sys
+import threading
 from asyncio import Queue
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import redirect_stdout, suppress
@@ -269,6 +270,61 @@ def wrapper(*args_, **kwargs_):
         return decorator(func_=func)
 
 
+class ThreadRoutedStdout:
+    """A stdout stand-in that routes each worker thread's writes to its own buffer
+
+    Threads that have not registered a buffer, the calling thread and tqdm among
+    them, write straight through to the wrapped stream. Installed once for the
+    duration of a parallelize() call, so unlike contextlib.redirect_stdout it
+    stays correct while several tasks print at the same time.
+    """
+
+    def __init__(self, stream: IO):
+        self.stream = stream
+        self._local = threading.local()
+
+    def register(self) -> List[str]:
+        """Gives the calling thread its own buffer and returns it"""
+        buffer = []
+        self._local.buffer = buffer
+        return buffer
+
+    def unregister(self):
+        """Sends the calling thread's writes back to the wrapped stream"""
+        self._local.buffer = None
+
+    def write(self, message):
+        buffer = getattr(self._local, "buffer", None)
+        if buffer is None:
+            return self.stream.write(message)
+        buffer.append(message)
+        return len(message)
+
+    def flush(self):
+        if getattr(self._local, "buffer", None) is None:
+            self.stream.flush()
+
+    def __getattr__(self, name):
+        # encoding, isatty, fileno and the like stay those of the wrapped stream
+        return getattr(self.stream, name)
+
+
+def capture_output(
+    call: Callable, stdout: ThreadRoutedStdout, sink: List[str], index: int
+) -> Callable:
+    """Wraps call so that what it prints ends up in sink[index] instead of stdout"""
+
+    def wrapper():
+        buffer = stdout.register()
+        try:
+            return call()
+        finally:
+            stdout.unregister()
+            sink[index] = "".join(buffer)
+
+    return wrapper
+
+
 def parallelize(
     func: Callable,
     iterable: Iterable,
@@ -286,8 +342,10 @@ def parallelize(
     iterable:
         The iterable, who's items are to be passed as singles to the function.
     flush_at_end:
-        If True, the message buffer will be printed (flushed) to stdout after parallel
-        execution.
+        If True, what the tasks print is written to stdout in the order of
+        ``iterable`` once the batch has finished. If False, it is discarded.
+        Either way it is captured while the tasks run, so that concurrent
+        printing cannot garble the progress bar.
     progress_bar:
         If True, a progress bar will be displayed.
     return_exceptions:
@@ -299,54 +357,46 @@ def parallelize(
         Keyword arguments to be passed to the function.z
     """
 
-    MODE = "asyncio"
-    if MODE == "dask":
-        # this mode can be removed if asyncio works as expected
-        with MessageBuffer(flush_at_end) as msg_buf:
-            tasks = [
-                dask.delayed(redirect_print_explicitly(func, msg_buf))(item, **kwargs)
-                for item in iterable
-            ]
-            print(
-                f"Performing parallel execution of {func.__name__} ({len(tasks)} tasks)."
+    items = list(iterable)
+    # One slot per item, so the collected output can be replayed in input order
+    # even though the tasks finish in whatever order they finish.
+    outputs = [""] * len(items)
+    stdout = ThreadRoutedStdout(sys.stdout)
+
+    # run func concurrently with each item in iterable using asyncio
+    async def _run_tasks():
+        loop = asyncio.get_event_loop()
+        # ToDo: handle IO-bound different from CPU-bound tasks
+        # see https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.run_in_executor
+        # None => ThreadPoolExecutor for IO-bound tasks
+        tasks = [
+            loop.run_in_executor(
+                None,
+                capture_output(
+                    functools.partial(func, item, **kwargs), stdout, outputs, index
+                ),
             )
-            if progress_bar:
-                with ProgressBar():
-                    results = dask.compute(*tasks)
-            else:
-                results = dask.compute(*tasks)
+            for index, item in enumerate(items)
+        ]
+        print(f"Performing parallel execution of {func.__name__} ({len(tasks)} tasks).")
+        if return_exceptions:
+            # tqdm.gather() re-raises the first exception (it has no
+            # return_exceptions kwarg), which would abandon the remaining
+            # results. Capture each task's exception instead so the batch
+            # completes and results stay aligned with the input order.
+            async def _capture(task):
+                try:
+                    return await task
+                except Exception as exc:
+                    return exc
 
-    if MODE == "asyncio":
-        # run func concurrently with each item in iterable using asyncio
-        async def _run_tasks():
-            loop = asyncio.get_event_loop()
-            # ToDo: handle IO-bound different from CPU-bound tasks
-            # see https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.run_in_executor
-            # None => ThreadPoolExecutor for IO-bound tasks
-            tasks = [
-                loop.run_in_executor(None, functools.partial(func, item, **kwargs))
-                for item in iterable
-            ]
-            print(
-                f"Performing parallel execution of {func.__name__} ({len(tasks)} tasks)."
-            )
-            if return_exceptions:
-                # tqdm.gather() re-raises the first exception (it has no
-                # return_exceptions kwarg), which would abandon the remaining
-                # results. Capture each task's exception instead so the batch
-                # completes and results stay aligned with the input order.
-                async def _capture(task):
-                    try:
-                        return await task
-                    except Exception as exc:
-                        return exc
-
-                tasks = [_capture(task) for task in tasks]
-            res = await tqdm.gather(
-                *tasks
-            )  # like asyncio.gather, but with a progress bar
-            return res
+            tasks = [_capture(task) for task in tasks]
+        # tqdm.gather is asyncio.gather plus a progress bar
+        gather = tqdm.gather if progress_bar else asyncio.gather
+        return await gather(*tasks)
 
+    sys.stdout = stdout
+    try:
         # Check if we are already in an event loop, e.g. in a Jupyter notebook
         try:
             loop = asyncio.get_running_loop()
@@ -362,6 +412,14 @@ async def _capture(task):
         else:
             # Starting new event loop
             results = asyncio.run(_run_tasks())
+    finally:
+        sys.stdout = stdout.stream
+        if flush_at_end:
+            # in the finally block, so a failed batch still reports what its
+            # tasks had to say
+            for output in outputs:
+                if output:
+                    print(output, end="")
 
     return results
 

tests/utils/parallelize_test.py

@@ -4,14 +4,18 @@
 kwargs forwarding, and the two failure modes: abort on the first exception, or
 collect exceptions in place so the batch completes with results still aligned
 to the input.
+
+Also covers issue #154: flush_at_end and progress_bar were dropped when the
+implementation moved from dask to asyncio, leaving both parameters inert.
 """
 
 import asyncio
+import sys
 import threading
 
 import pytest
 
-from osw.utils.util import parallelize
+from osw.utils.util import ThreadRoutedStdout, parallelize
 
 # Generous enough that a loaded CI machine still passes, short enough that a
 # serialized execution trips it instead of hanging the suite.
@@ -112,3 +116,165 @@ async def caller():
         return parallelize(_double, [1, 2, 3])
 
     assert asyncio.run(caller()) == [2, 4, 6]
+
+
+def _print_then_return(item):
+    print(f"handled {item}")
+    return item
+
+
+def test_flush_at_end_prints_what_the_tasks_printed(capsys):
+    parallelize(_print_then_return, [1, 2, 3], flush_at_end=True)
+
+    out = capsys.readouterr().out
+    assert "handled 1" in out
+    assert "handled 2" in out
+    assert "handled 3" in out
+
+
+def test_flush_at_end_replays_in_input_order(capsys):
+    """Output follows the iterable, not the order the tasks happened to finish."""
+    import time
+
+    delays = {1: 0.05, 2: 0.0, 3: 0.03}
+
+    def sleep_then_print(item):
+        time.sleep(delays[item])
+        print(f"handled {item}")
+
+    parallelize(sleep_then_print, [1, 2, 3], flush_at_end=True)
+
+    out = capsys.readouterr().out
+    assert out.index("handled 1") < out.index("handled 2") < out.index("handled 3")
+
+
+def test_flush_at_end_defers_until_the_batch_is_done():
+    """Nothing reaches the real stdout while the tasks are still running."""
+    written = []
+    snapshots = []
+    barrier = threading.Barrier(3, timeout=BARRIER_TIMEOUT)
+
+    class _Recorder:
+        def write(self, message):
+            written.append(message)
+            return len(message)
+
+        def flush(self):
+            pass
+
+    def print_then_look(item):
+        print(f"handled {item}")
+        barrier.wait()  # every task has printed by the time this releases
+        snapshots.append("".join(written))
+
+    original = sys.stdout
+    sys.stdout = _Recorder()
+    try:
+        parallelize(print_then_look, [1, 2, 3], flush_at_end=True, progress_bar=False)
+    finally:
+        sys.stdout = original
+
+    # mid-flight the real stream had seen none of it, afterwards it has all of it
+    assert not any("handled" in snapshot for snapshot in snapshots)
+    assert "".join(written).count("handled") == 3
+
+
+def test_task_output_is_suppressed_without_flush_at_end(capsys):
+    parallelize(_print_then_return, [1, 2, 3])
+
+    assert "handled" not in capsys.readouterr().out
+
+
+def test_flush_at_end_reports_even_when_the_batch_fails(capsys):
+    """A failing item must not swallow the diagnostics of the others."""
+
+    def print_then_fail(item):
+        print(f"handled {item}")
+        if item == 2:
+            raise ValueError("boom")
+        return item
+
+    with pytest.raises(ValueError, match="boom"):
+        parallelize(print_then_fail, [1, 2, 3], flush_at_end=True)
+
+    assert "handled 1" in capsys.readouterr().out
+
+
+def test_stdout_is_restored_afterwards():
+    original = sys.stdout
+
+    parallelize(_print_then_return, [1, 2], flush_at_end=True)
+
+    assert sys.stdout is original
+
+
+def test_stdout_is_restored_after_a_failure():
+    original = sys.stdout
+
+    with pytest.raises(ValueError):
+        parallelize(_fail_on_even, [2])
+
+    assert sys.stdout is original
+
+
+def test_progress_bar_can_be_switched_off(capsys):
+    assert parallelize(_double, [1, 2, 3], progress_bar=False) == [2, 4, 6]
+
+    # tqdm renders to stderr, so switching it off must leave stderr untouched
+    assert capsys.readouterr().err == ""
+
+
+def test_progress_bar_is_shown_by_default(capsys):
+    assert parallelize(_double, [1, 2, 3]) == [2, 4, 6]
+
+    assert "3/3" in capsys.readouterr().err
+
+
+def test_both_parameters_combine():
+    """Neither parameter may break the results the caller is actually after."""
+    assert parallelize(
+        _print_then_return, [1, 2, 3], flush_at_end=True, progress_bar=False
+    ) == [1, 2, 3]
+
+
+def test_routed_stdout_keeps_threads_apart():
+    """The property the old redirect_stdout lacked: concurrent writes stay separate."""
+    stdout = ThreadRoutedStdout(sys.stdout)
+    collected = {}
+    barrier = threading.Barrier(2, timeout=BARRIER_TIMEOUT)
+
+    def write_interleaved(name):
+        buffer = stdout.register()
+        barrier.wait()  # force both threads to be inside their capture at once
+        stdout.write(name)
+        barrier.wait()
+        stdout.write(name)
+        stdout.unregister()
+        collected[name] = "".join(buffer)
+
+    threads = [
+        threading.Thread(target=write_interleaved, args=(name,)) for name in ("a", "b")
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert collected == {"a": "aa", "b": "bb"}
+
+
+def test_routed_stdout_passes_unregistered_threads_through():
+    class _Sink:
+        def __init__(self):
+            self.written = []
+
+        def write(self, message):
+            self.written.append(message)
+            return len(message)
+
+    sink = _Sink()
+    stdout = ThreadRoutedStdout(sink)
+
+    stdout.write("straight through")
+
+    assert sink.written == ["straight through"]