test_threaded_redirect.py

test_threaded_redirect.py (764 bytes)

import concurrent.futures
import sys
import io
import time
import threaded_redirect

def print_slowly(message, count=3, delay=0.1):
    for _ in range(count):
        print(message)
        time.sleep(delay)

def check_print_slowly(message):
    sio = io.StringIO()
    with threaded_redirect.redirect_stdout(sio):
        print_slowly(message)
    assert sio.getvalue() == 3 * f"{message}\n"
    print(f"OK {message}")

orig_stdout = sys.stdout
check_print_slowly(3)
assert orig_stdout is sys.stdout
check_print_slowly(4)
assert orig_stdout is sys.stdout
check_print_slowly(5)
assert orig_stdout is sys.stdout

with concurrent.futures.ThreadPoolExecutor() as executor:
    for i in range(4):
        executor.submit(check_print_slowly, i)
    executor.shutdown()

threaded_redirect.py

threaded_redirect.py (1840 bytes)

import sys
import contextlib
import typing
import threading
import collections


class _MultiStream:
    def __init__(
        self,
        map: typing.Mapping[threading.Thread, list[typing.TextIO]],
        default: typing.TextIO,
    ) -> None:
        self._map = map
        self._default = [default]

    def __getattr__(self, attr):
        sseq = self._map.get(threading.current_thread(), []) or self._default
        stream = sseq[0]
        return getattr(stream, attr)


class _IoReplacer:
    def __init__(self, stream_name: str) -> None:
        self.stream_name = stream_name
        self.map: collections.defaultdict[threading.Thread, list[typing.TextIO]] = (
            collections.defaultdict(list)
        )
        self.old_stream: typing.TextIO = getattr(sys, self.stream_name)
        self.lock = threading.Lock()
        self.stream = _MultiStream(self.map, self.old_stream)

    def push(self, target: typing.TextIO) -> None:
        thread = threading.current_thread()
        with self.lock:
            if not self.map:
                setattr(sys, self.stream_name, self.stream)
            self.map[thread].append(target)

    def pop(self) -> None:
        thread = threading.current_thread()
        with self.lock:
            self.map[thread].pop()
            if not self.map[thread]:
                del self.map[thread]
            if not self.map:
                setattr(sys, self.stream_name, self.old_stream)

    @contextlib.contextmanager
    def __call__(self, new_target: typing.TextIO) -> typing.Iterator[None]:
        try:
            self.push(new_target)
            yield
        finally:
            assert self.map[threading.current_thread()][-1] is new_target
            self.pop()
            self.push


redirect_stdout = _IoReplacer("stdout")
redirect_stderr = _IoReplacer("stderr")