use std::collections::HashMap;
use std::io::{Write as _, stderr};
use std::process::Stdio;
use color_print::cformat;
use worktrunk::config::{UserConfig, expand_template};
use worktrunk::git::{Repository, WorktreeInfo, WorktrunkError, interrupt_exit_code};
use worktrunk::shell_exec::Cmd;
use worktrunk::styling::{
eprintln, error_message, format_with_gutter, progress_message, success_message, warning_message,
};
use crate::commands::command_executor::{CommandContext, build_hook_context};
use crate::commands::worktree_display_name;
pub fn step_for_each(args: Vec<String>, format: crate::cli::SwitchFormat) -> anyhow::Result<()> {
let json_mode = format == crate::cli::SwitchFormat::Json;
let repo = Repository::current()?;
let worktrees: Vec<&WorktreeInfo> = repo
.list_worktrees()?
.iter()
.filter(|wt| !wt.is_prunable())
.collect();
let config = UserConfig::load()?;
let mut failed: Vec<String> = Vec::new();
let mut json_results: Vec<serde_json::Value> = Vec::new();
let mut interrupted: Option<i32> = None;
let total = worktrees.len();
for &wt in &worktrees {
let display_name = worktree_display_name(wt, &repo, &config);
eprintln!(
"{}",
progress_message(format!("Running in {display_name}..."))
);
let ctx = CommandContext::new(&repo, &config, wt.branch.as_deref(), &wt.path, false);
let context_map = build_hook_context(&ctx, &[])?;
let vars: HashMap<&str, &str> = context_map
.iter()
.map(|(k, v)| (k.as_str(), v.as_str()))
.collect();
let expanded: Vec<String> = args
.iter()
.map(|arg| expand_template(arg, &vars, false, &repo, "for-each argument"))
.collect::<Result<_, _>>()?;
let context_json = serde_json::to_string(&context_map)
.expect("HashMap<String, String> serialization should never fail");
match run_argv(&wt.path, expanded, &context_json) {
Ok(()) => {
if json_mode {
json_results.push(serde_json::json!({
"branch": wt.branch,
"path": wt.path,
"exit_code": 0,
"success": true,
}));
}
}
Err(err) => {
let signal_exit = interrupt_exit_code(&err);
let (exit_info, exit_code, error_msg, show_detail) =
if let Some(WorktrunkError::ChildProcessExited { code, message, .. }) =
err.downcast_ref::<WorktrunkError>()
{
(
format!(" (exit code {code})"),
serde_json::json!(code),
message.clone(),
false,
)
} else {
let msg = err.to_string();
(
" (spawn failed)".to_string(),
serde_json::json!(null),
msg,
true,
)
};
eprintln!(
"{}",
error_message(cformat!("Failed in <bold>{display_name}</>{exit_info}"))
);
if show_detail {
eprintln!("{}", format_with_gutter(&error_msg, None));
}
failed.push(display_name.to_string());
if json_mode {
json_results.push(serde_json::json!({
"branch": wt.branch,
"path": wt.path,
"exit_code": exit_code,
"success": false,
"error": error_msg,
}));
}
if let Some(code) = signal_exit {
interrupted = Some(code);
break;
}
}
}
}
if let Some(exit_code) = interrupted {
if json_mode {
println!("{}", serde_json::to_string_pretty(&json_results)?);
} else {
eprintln!();
eprintln!(
"{}",
warning_message("Interrupted — skipped remaining worktrees")
);
}
return Err(WorktrunkError::AlreadyDisplayed { exit_code }.into());
}
if json_mode {
println!("{}", serde_json::to_string_pretty(&json_results)?);
if failed.is_empty() {
return Ok(());
} else {
return Err(WorktrunkError::AlreadyDisplayed { exit_code: 1 }.into());
}
}
eprintln!();
if failed.is_empty() {
eprintln!(
"{}",
success_message(format!(
"Completed in {total} worktree{}",
if total == 1 { "" } else { "s" }
))
);
Ok(())
} else {
eprintln!(
"{}",
warning_message(format!(
"{} of {total} worktree{} failed",
failed.len(),
if total == 1 { "" } else { "s" }
))
);
let failed_list = failed.join("\n");
eprintln!("{}", format_with_gutter(&failed_list, None));
Err(WorktrunkError::AlreadyDisplayed { exit_code: 1 }.into())
}
}
fn run_argv(
working_dir: &std::path::Path,
argv: Vec<String>,
stdin_json: &str,
) -> anyhow::Result<()> {
stderr().flush()?;
eprint!("{}", anstyle::Reset);
stderr().flush().ok();
let mut iter = argv.into_iter();
let program = iter
.next()
.expect("clap enforces at least one argv element");
Cmd::new(program)
.args(iter)
.current_dir(working_dir)
.stdout(Stdio::from(std::io::stderr()))
.forward_signals()
.stdin_bytes(stdin_json.as_bytes().to_vec())
.stream()?;
stderr().flush()?;
Ok(())
}