rp/pr/multiplexer.py

import threading
import queue
import time
import sys
from pr.ui import Colors

class TerminalMultiplexer:
    def __init__(self, name, show_output=True):
        self.name = name
        self.show_output = show_output
        self.stdout_buffer = []
        self.stderr_buffer = []
        self.stdout_queue = queue.Queue()
        self.stderr_queue = queue.Queue()
        self.active = True
        self.lock = threading.Lock()

        if self.show_output:
            self.display_thread = threading.Thread(target=self._display_worker, daemon=True)
            self.display_thread.start()

    def _display_worker(self):
        while self.active:
            try:
                line = self.stdout_queue.get(timeout=0.1)
                if line:
                    sys.stdout.write(f"{Colors.GRAY}[{self.name}]{Colors.RESET} {line}")
                    sys.stdout.flush()
            except queue.Empty:
                pass

            try:
                line = self.stderr_queue.get(timeout=0.1)
                if line:
                    sys.stderr.write(f"{Colors.YELLOW}[{self.name} err]{Colors.RESET} {line}")
                    sys.stderr.flush()
            except queue.Empty:
                pass

    def write_stdout(self, data):
        with self.lock:
            self.stdout_buffer.append(data)
        if self.show_output:
            self.stdout_queue.put(data)

    def write_stderr(self, data):
        with self.lock:
            self.stderr_buffer.append(data)
        if self.show_output:
            self.stderr_queue.put(data)

    def get_stdout(self):
        with self.lock:
            return ''.join(self.stdout_buffer)

    def get_stderr(self):
        with self.lock:
            return ''.join(self.stderr_buffer)

    def get_all_output(self):
        with self.lock:
            return {
                'stdout': ''.join(self.stdout_buffer),
                'stderr': ''.join(self.stderr_buffer)
            }

    def close(self):
        self.active = False
        if hasattr(self, 'display_thread'):
            self.display_thread.join(timeout=1)

_multiplexers = {}
_mux_counter = 0
_mux_lock = threading.Lock()

def create_multiplexer(name=None, show_output=True):
    global _mux_counter
    with _mux_lock:
        if name is None:
            _mux_counter += 1
            name = f"process-{_mux_counter}"
        mux = TerminalMultiplexer(name, show_output)
        _multiplexers[name] = mux
        return name, mux

def get_multiplexer(name):
    return _multiplexers.get(name)

def close_multiplexer(name):
    mux = _multiplexers.get(name)
    if mux:
        mux.close()
        del _multiplexers[name]

def cleanup_all_multiplexers():
    for mux in list(_multiplexers.values()):
        mux.close()
    _multiplexers.clear()