use miette::{IntoDiagnostic, WrapErr, miette};
use std::io::IsTerminal;
use tokio::io::{AsyncBufReadExt, AsyncRead, BufReader};
use tokio::process::Command;
const PALETTE: &[u8] = &[31, 32, 33, 34, 35, 36];
#[derive(Debug, Clone)]
pub(crate) enum OutputMode {
Prefix { name: String, color_index: usize },
NoPrefix,
}
impl OutputMode {
pub(crate) fn prefix(name: Option<&str>, index: usize) -> Self {
match name {
Some(n) => Self::Prefix {
name: n.to_string(),
color_index: index % PALETTE.len(),
},
None => Self::NoPrefix,
}
}
}
pub(crate) async fn run_command(
mut cmd: Command,
mode: &OutputMode,
) -> miette::Result<std::process::ExitStatus> {
cmd.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
let mut child = cmd
.spawn()
.into_diagnostic()
.wrap_err("failed to spawn child process")?;
let stdout = child
.stdout
.take()
.expect("stdout was piped above; take() must succeed");
let stderr = child
.stderr
.take()
.expect("stderr was piped above; take() must succeed");
let no_color = std::env::var_os("NO_COLOR").is_some();
let stdout_prefix = format_line_prefix(mode, std::io::stdout().is_terminal() && !no_color);
let stderr_prefix = format_line_prefix(mode, std::io::stderr().is_terminal() && !no_color);
let stdout_pump = tokio::spawn(pump_lines(stdout, stdout_prefix, false));
let stderr_pump = tokio::spawn(pump_lines(stderr, stderr_prefix, true));
let status = child
.wait()
.await
.into_diagnostic()
.wrap_err("failed to wait on child process")?;
stdout_pump
.await
.map_err(|e| miette!("stdout pump task failed: {e}"))?
.map_err(|e| miette!("stdout pump io error: {e}"))?;
stderr_pump
.await
.map_err(|e| miette!("stderr pump task failed: {e}"))?
.map_err(|e| miette!("stderr pump io error: {e}"))?;
Ok(status)
}
fn format_line_prefix(mode: &OutputMode, color: bool) -> String {
match mode {
OutputMode::NoPrefix => String::new(),
OutputMode::Prefix { name, color_index } => {
if color {
let code = PALETTE[*color_index];
format!("\x1b[{code}m{name}\x1b[0m: ")
} else {
format!("{name}: ")
}
}
}
}
async fn pump_lines<R>(reader: R, prefix: String, is_stderr: bool) -> std::io::Result<()>
where
R: AsyncRead + Unpin,
{
let mut reader = BufReader::new(reader);
let mut line = String::new();
loop {
line.clear();
let n = reader.read_line(&mut line).await?;
if n == 0 {
return Ok(());
}
let trimmed = line.trim_end_matches(['\n', '\r']);
if is_stderr {
aube_scripts::write_line_to_real_stderr(&format!("{prefix}{trimmed}"));
} else {
println!("{prefix}{trimmed}");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn run_command_with_prefix_drains_both_streams() {
if cfg!(windows) {
return;
}
let mut cmd = tokio::process::Command::new("sh");
cmd.arg("-c").arg("echo out1; echo err1 1>&2; echo out2");
let mode = OutputMode::Prefix {
name: "demo".to_string(),
color_index: 0,
};
let status = run_command(cmd, &mode).await.unwrap();
assert!(status.success());
}
#[tokio::test]
async fn run_command_no_prefix_mode_still_pipes() {
let mut cmd = tokio::process::Command::new("sh");
cmd.arg("-c").arg("true");
let status = run_command(cmd, &OutputMode::NoPrefix).await.unwrap();
assert!(status.success());
}
#[test]
fn prefix_color_rotates_modulo_palette() {
let a = OutputMode::prefix(Some("foo"), 0);
let b = OutputMode::prefix(Some("foo"), PALETTE.len());
match (a, b) {
(
OutputMode::Prefix { color_index: i, .. },
OutputMode::Prefix { color_index: j, .. },
) => assert_eq!(i, j),
_ => panic!("named pkg should produce Prefix mode"),
}
}
#[test]
fn unnamed_pkg_falls_back_to_noprefix() {
assert!(matches!(OutputMode::prefix(None, 3), OutputMode::NoPrefix));
}
#[test]
fn format_line_prefix_skips_color_when_disabled() {
let p = format_line_prefix(
&OutputMode::Prefix {
name: "demo".to_string(),
color_index: 0,
},
false,
);
assert_eq!(p, "demo: ");
}
#[test]
fn format_line_prefix_emits_ansi_when_color_enabled() {
let p = format_line_prefix(
&OutputMode::Prefix {
name: "demo".to_string(),
color_index: 0,
},
true,
);
assert_eq!(p, "\x1b[31mdemo\x1b[0m: ");
}
}