keep color information from pty terminal

.github/workflows/tests.yaml

@@ -19,6 +19,7 @@ jobs:
           - "3.10"
         os:
           - ubuntu-latest
+          - macos-latest
 
     steps:
       - uses: actions/checkout@v2

app/package.json

@@ -11,6 +11,7 @@
 	},
 	"dependencies": {
 		"@xyflow/react": "^12.4.2",
+		"ansi-to-react": "^6.1.6",
 		"bootstrap": "^5.3.3",
 		"elkjs": "^0.9.3",
 		"react": "^18.3.1",

app/src/GraphStateNode.tsx

@@ -7,6 +7,7 @@ import Modal from "react-bootstrap/Modal";
 import Markdown from "react-markdown";
 import { GraphNode } from "./types";
 import GraphContext from "./GraphContext";
+import Ansi from "ansi-to-react";
 
 interface GraphStateNodeProps {
 	data: {
@@ -159,13 +160,25 @@ ${data.node.deps_hash}
 					{nodeData.stdout && (
 						<>
 							<h5>STDOUT</h5>
-							<pre>{nodeData.stdout}</pre>
+							<div style={{ overflowX: 'auto' }}>
+							{nodeData.stdout.split('\n').map((line, index) => (
+								<><Ansi key={index}>{line}</Ansi> <br /></>
+
+								))}
+								</div>
 						</>
 					)}
 					{nodeData.stderr && (
 						<>
 							<h5>STDERR</h5>
-							<pre>{nodeData.stderr}</pre>
+							<div style={{ overflowX: 'auto' }}>
+							{nodeData.stderr.split('\n').map((line, index) => (
+								<>
+									<Ansi key={index}>{line}</Ansi>
+									<br />
+								</>
+								))}
+								</div>
 						</>
 					)}
 					{nodeData.worker && (

paraffin/stage.py

@@ -3,12 +3,13 @@
 import dataclasses
 import json
 import logging
+import re
 import subprocess
-import threading
 import time
 from pathlib import Path
 
 import dvc.api
+import pexpect
 import yaml
 from dvc.lock import LockError
 from dvc.stage import PipelineStage
@@ -20,6 +21,9 @@
 log = logging.getLogger(__name__)
 
 
+ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
+
 @dataclasses.dataclass(frozen=True, eq=True)
 class PipelineStageDC:
     """Container for a DVC stage."""
@@ -83,54 +87,35 @@ def get_lock(name: str) -> tuple[dict, str]:
     return stage_lock, deps_hash
 
 
-def _stream_reader(pipe, callback) -> None:
-    """Reads lines from a pipe and calls the callback function."""
-    with pipe:
-        for line in iter(pipe.readline, ""):  # Read until EOF
-            callback(line)
-
-
 def run_command(command: list[str]) -> tuple[int, str, str]:
-    """Run a subprocess command, capturing its stdout, stderr, and return code."""
-    process = subprocess.Popen(
-        command,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        universal_newlines=True,
-    )
+    """Runs a command with PTY support, capturing stdout/stderr while preserving colors and progress bars."""
+
+    cmd_str = " ".join(command)  # Convert list to string
+    child = pexpect.spawn(cmd_str, encoding="utf-8", timeout=None)
 
     stdout_lines = []
     stderr_lines = []
 
-    def print_and_store_stdout(line):
-        print(line, end="")  # Print in real-time
-        stdout_lines.append(line)
+    # Capture output in real-time
+    try:
+        while True:
+            line = child.readline()
+            if not line:
+                break  # Stop when no more output
 
-    def print_and_store_stderr(line):
-        print(line, end="")  # Print in real-time
-        stderr_lines.append(line)
+            print(line, end="")  # Print output live (preserves colors)
+            stdout_lines.append(ANSI_ESCAPE.sub("", line))  # Capture stdout
 
-    # Create threads to read stdout and stderr
-    stdout_thread = threading.Thread(
-        target=_stream_reader,
-        args=(process.stdout, print_and_store_stdout),
-        daemon=True,
-    )
-    stderr_thread = threading.Thread(
-        target=_stream_reader,
-        args=(process.stderr, print_and_store_stderr),
-        daemon=True,
-    )
-
-    stdout_thread.start()
-    stderr_thread.start()
+    except pexpect.EOF:
+        pass  # Process finished
 
-    return_code = process.wait()  # Ensure process completes
+    return_code = child.wait()
 
-    stdout_thread.join()
-    stderr_thread.join()
-
-    return return_code, "".join(stdout_lines), "".join(stderr_lines)
+    return (
+        return_code,
+        "".join(stdout_lines),
+        "".join(stderr_lines),
+    )  # stderr is empty (pexpect merges output)
 
 
 @retry(10, (LockError,), delay=0.5)

pyproject.toml

@@ -11,6 +11,7 @@ requires-python = ">=3.10"
 dependencies = [
     "dvc>=3.59.0",
     "fastapi[standard]>=0.115.7",
+    "pexpect>=4.9.0",
     "sqlmodel>=0.0.22",
     "typer>=0.13",
 ]