|
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()
|