use std::fmt::Write;
use clap::Command;
use console::style;
pub struct TreeBuilder {
output: String,
indent_stack: Vec<bool>,
max_item_width: usize,
}
impl TreeBuilder {
fn new() -> Self {
Self {
output: String::with_capacity(4096),
indent_stack: Vec::with_capacity(8),
max_item_width: 0,
}
}
fn build(mut self, cmd: &Command, root_label: &str) -> String {
self.measure_widths(cmd, 0);
self.max_item_width += 2;
writeln!(&mut self.output, "{}", style(root_label).yellow().bold()).unwrap();
self.render_command(cmd);
self.output
}
fn measure_widths(&mut self, cmd: &Command, depth: usize) {
let indent = depth * 4; let branch = 4;
for arg in cmd.get_arguments() {
if Self::skip_arg(arg) {
continue;
}
let w = indent + branch + Self::arg_display_width(arg);
self.max_item_width = self.max_item_width.max(w);
}
for sub in cmd.get_subcommands() {
if sub.is_hide_set() {
continue;
}
let w = indent + branch + Self::subcommand_display_width(sub);
self.max_item_width = self.max_item_width.max(w);
self.measure_widths(sub, depth + 1);
}
}
fn arg_display_width(arg: &clap::Arg) -> usize {
let id = arg.get_id().as_str();
if arg.get_short().is_none() && arg.get_long().is_none() {
let name = arg
.get_value_names()
.and_then(|v| v.first().map(|n| n.as_str()))
.unwrap_or(id);
if arg.is_last_set() {
return "[-- ".len() + name.to_uppercase().len() + "...]".len();
}
return name.to_uppercase().len() + 2; }
let mut w = 0usize;
if let Some(_short) = arg.get_short() {
w += 2; if arg.get_long().is_some() {
w += 2; }
}
if let Some(long) = arg.get_long() {
w += 2 + long.len(); }
if arg.get_num_args().is_some() || arg.get_action().takes_values() {
if let Some(names) = arg.get_value_names() {
if let Some(name) = names.first() {
w += 1 + name.to_uppercase().len() + 2; }
} else {
w += " <VALUE>".len();
}
}
let short_aliases: Vec<_> = arg.get_visible_short_aliases().unwrap_or_default();
let long_aliases: Vec<_> = arg.get_visible_aliases().unwrap_or_default();
if !short_aliases.is_empty() || !long_aliases.is_empty() {
w += " (aliases: ".len();
let mut first = true;
for _ in &short_aliases {
if !first {
w += 2;
} w += 2; first = false;
}
for a in &long_aliases {
if !first {
w += 2;
}
w += 2 + a.len(); first = false;
}
w += 1; }
w
}
fn subcommand_display_width(cmd: &Command) -> usize {
let mut w = cmd.get_name().len();
let aliases: Vec<_> = cmd.get_visible_aliases().collect();
if !aliases.is_empty() {
w += " (aliases: ".len();
w += aliases.iter().map(|a| a.len()).sum::<usize>();
w += (aliases.len() - 1) * 2; w += 1; }
w
}
fn render_command(&mut self, cmd: &Command) {
let mut positionals: Vec<&clap::Arg> = Vec::new();
let mut flags: Vec<&clap::Arg> = Vec::new();
for arg in cmd.get_arguments() {
if Self::skip_arg(arg) {
continue;
}
if arg.get_short().is_none() && arg.get_long().is_none() {
positionals.push(arg);
} else {
flags.push(arg);
}
}
let subcommands: Vec<&Command> =
cmd.get_subcommands().filter(|s| !s.is_hide_set()).collect();
let total = positionals.len() + flags.len() + subcommands.len();
let mut idx = 0;
for arg in &positionals {
idx += 1;
self.render_arg(arg, idx == total);
}
for arg in &flags {
idx += 1;
self.render_arg(arg, idx == total);
}
for sub in &subcommands {
idx += 1;
self.render_subcommand(sub, idx == total);
}
}
fn render_arg(&mut self, arg: &clap::Arg, is_last: bool) {
let label = Self::format_arg(arg);
let description = arg.get_help().map(|h| {
let s = h.to_string();
s.lines().next().unwrap_or("").to_string()
});
let prefix = self.build_prefix(is_last);
let styled_label = Self::style_arg(&label);
let current_indent = self.indent_stack.len() * 4;
let total_width = current_indent + 4 + label.len(); let pad = self.max_item_width.saturating_sub(total_width);
write!(&mut self.output, "{}{}", style(&prefix).dim(), styled_label).unwrap();
if let Some(desc) = description.filter(|d| !d.is_empty()) {
write!(&mut self.output, "{:width$}{}", "", desc, width = pad).unwrap();
}
writeln!(&mut self.output).unwrap();
}
fn render_subcommand(&mut self, cmd: &Command, is_last: bool) {
let name = cmd.get_name();
let aliases: Vec<_> = cmd.get_visible_aliases().collect();
let label_width = Self::subcommand_display_width(cmd);
let current_indent = self.indent_stack.len() * 4;
let total_width = current_indent + 4 + label_width;
let pad = self.max_item_width.saturating_sub(total_width);
let prefix = self.build_prefix(is_last);
let colored_name = match self.indent_stack.len() {
0 => style(name).magenta().bold().to_string(),
1 => style(name).blue().bold().to_string(),
2 => style(name).green().bold().to_string(),
_ => style(name).cyan().bold().to_string(),
};
write!(&mut self.output, "{}{}", style(&prefix).dim(), colored_name).unwrap();
if !aliases.is_empty() {
write!(
&mut self.output,
" {}",
style(format!("(aliases: {})", aliases.join(", "))).dim()
)
.unwrap();
}
if let Some(about) = cmd.get_about() {
let about_str = about.to_string();
if let Some(line) = about_str.lines().next().filter(|l| !l.is_empty()) {
write!(&mut self.output, "{:width$}{}", "", line, width = pad).unwrap();
}
}
writeln!(&mut self.output).unwrap();
let has_children = cmd.get_subcommands().any(|s| !s.is_hide_set())
|| cmd.get_arguments().any(|a| !Self::skip_arg(a));
if has_children {
self.indent_stack.push(!is_last);
self.render_command(cmd);
self.indent_stack.pop();
}
}
fn build_prefix(&self, is_last: bool) -> String {
let depth = self.indent_stack.len();
let mut prefix = String::with_capacity(depth * 4 + 4);
for &continues in &self.indent_stack {
if continues {
prefix.push_str("│ ");
} else {
prefix.push_str(" ");
}
}
if is_last {
prefix.push_str("└── ");
} else {
prefix.push_str("├── ");
}
prefix
}
fn format_arg(arg: &clap::Arg) -> String {
let id = arg.get_id().as_str();
if arg.get_short().is_none() && arg.get_long().is_none() {
let name = arg
.get_value_names()
.and_then(|v| v.first().map(|n| n.as_str()))
.unwrap_or(id);
if arg.is_last_set() {
return format!("[-- {}...]", name.to_uppercase());
}
return if arg.is_required_set() {
format!("<{}>", name.to_uppercase())
} else {
format!("[{}]", name.to_uppercase())
};
}
let mut s = String::with_capacity(32);
if let Some(short) = arg.get_short() {
write!(&mut s, "-{}", short).unwrap();
if arg.get_long().is_some() {
s.push_str(", ");
}
}
if let Some(long) = arg.get_long() {
write!(&mut s, "--{}", long).unwrap();
}
if arg.get_num_args().is_some() || arg.get_action().takes_values() {
if let Some(names) = arg.get_value_names() {
if let Some(name) = names.first() {
write!(&mut s, " <{}>", name.to_uppercase()).unwrap();
}
} else {
s.push_str(" <VALUE>");
}
}
let short_aliases: Vec<_> = arg.get_visible_short_aliases().unwrap_or_default();
let long_aliases: Vec<_> = arg.get_visible_aliases().unwrap_or_default();
if !short_aliases.is_empty() || !long_aliases.is_empty() {
s.push_str(" (aliases: ");
let mut first = true;
for a in &short_aliases {
if !first {
s.push_str(", ");
}
write!(&mut s, "-{}", a).unwrap();
first = false;
}
for a in &long_aliases {
if !first {
s.push_str(", ");
}
write!(&mut s, "--{}", a).unwrap();
first = false;
}
s.push(')');
}
s
}
fn style_arg(label: &str) -> String {
if label.starts_with('-') {
if let Some((flag_part, alias_part)) = label.split_once(" (aliases: ") {
let styled_flag = if let Some((fl, val)) = flag_part.split_once(' ') {
format!("{} {}", style(fl).dim(), style(val).dim())
} else {
style(flag_part).dim().to_string()
};
return format!(
"{} {}",
styled_flag,
style(format!("(aliases: {}", alias_part)).dim()
);
}
if let Some((fl, val)) = label.split_once(' ') {
return format!("{} {}", style(fl).dim(), style(val).dim());
}
return style(label).dim().to_string();
}
if (label.starts_with('<') && label.ends_with('>'))
|| (label.starts_with('[') && label.ends_with(']'))
|| label.starts_with("[-- ")
{
return style(label).dim().to_string();
}
label.to_string()
}
fn skip_arg(arg: &clap::Arg) -> bool {
let id = arg.get_id().as_str();
id == "help" || id == "version" || id == "tree" || arg.is_hide_set()
}
}
pub fn generate_tree(cmd: &Command) -> String {
TreeBuilder::new().build(cmd, cmd.get_name())
}
pub fn generate_tree_with_root(cmd: &Command, root: &str) -> String {
TreeBuilder::new().build(cmd, root)
}
pub fn try_show_tree(cmd: &Command) -> Option<String> {
let args: Vec<String> = std::env::args().collect();
if !args.iter().any(|a| a == "--tree") {
return None;
}
let (path, deepest) = find_deepest_subcommand(cmd, &args);
let tree = if path.len() > 1 {
generate_tree_with_root(&deepest, &path.join(" "))
} else {
generate_tree(&deepest)
};
Some(tree)
}
fn find_deepest_subcommand(cmd: &Command, args: &[String]) -> (Vec<String>, Command) {
let mut path = vec![cmd.get_name().to_string()];
let mut current = cmd.clone();
for arg in args.iter().skip(1) {
if arg.starts_with('-') {
continue;
}
if let Some(sub) = current.find_subcommand(arg) {
if sub.is_hide_set() {
continue;
}
path.push(arg.to_string());
current = sub.clone();
}
}
(path, current)
}