From eae04b10b64ac1411defab285e006b737e1be694 Mon Sep 17 00:00:00 2001 From: Basil Crow Date: Thu, 5 Mar 2026 10:01:20 -0800 Subject: [PATCH] ptree: Truncate output lines to terminal width Add a -w/--wrap flag to allow lines to wrap; by default, lines are truncated to the terminal width (via TIOCGWINSZ or ). Refactor cmd_summary_from() to return a String instead of printing directly, enabling callers to measure and truncate the line before output. Co-Authored-By: Claude Opus 4.6 --- Cargo.toml | 2 +- build.rs | 7 ++- src/bin/pstack.rs | 4 +- src/bin/ptree.rs | 117 +++++++++++++++++++++++++++----------------- src/display.rs | 26 +++++----- tests/ptree_test.rs | 90 ++++++++++++++++++++++++++++++++++ 6 files changed, 183 insertions(+), 63 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3737138..b45ee14 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,7 +20,7 @@ flate2 = "1.1" foreign-types = "0.5" lexopt = "0.3" libc = "0.2" -nix = { version = "0.30", features = ["event", "feature", "fs", "process", "resource", "sched", "signal", "socket", "time", "user"] } +nix = { version = "0.30", features = ["event", "feature", "fs", "ioctl", "process", "resource", "sched", "signal", "socket", "time", "user"] } rustc-demangle = "0.1" zstd = { version = "0.13", default-features = false } diff --git a/build.rs b/build.rs index ca7140a..bb27427 100644 --- a/build.rs +++ b/build.rs @@ -472,7 +472,7 @@ fn main() { child processes indented from their respective parent processes. An \ argument of all digits is taken to be a process ID; otherwise it is \ assumed to be a user login name. The default is all processes.", - synopsis: "[-ag] [pid|user]...", + synopsis: "[-agw] [pid|user]...", options: &[ ( "-a, --all", @@ -484,6 +484,11 @@ fn main() { locale, the UTF-8 line drawing characters are used, otherwise \ ASCII line drawing characters are used.", ), + ( + "-w, --wrap", + "Allow output lines to wrap. Normally output lines are \ + truncated to the current width of the terminal window.", + ), ], operands: &[ ( diff --git a/src/bin/pstack.rs b/src/bin/pstack.rs index 74d33f4..bdc1e6b 100644 --- a/src/bin/pstack.rs +++ b/src/bin/pstack.rs @@ -105,7 +105,7 @@ fn print_stack( handle, |pid| { print!("{pid}:\t"); - ptools::display::print_cmd_summary_from(handle); + println!("{}", ptools::display::cmd_summary_from(handle)); println!(); first_thread.set(true); }, @@ -116,7 +116,7 @@ fn print_stack( // attaches to (and unwinds) this one TID instead of every thread. opts.tid(tid as u32); print!("{tid}:\t"); - ptools::display::print_cmd_summary_from(handle); + println!("{}", ptools::display::cmd_summary_from(handle)); opts.trace_each(handle, |thread| { print_thread(&thread, Some(tid), args.max_frames, false, &first_thread); })?; diff --git a/src/bin/ptree.rs b/src/bin/ptree.rs index 80e149e..1b424e5 100644 --- a/src/bin/ptree.rs +++ b/src/bin/ptree.rs @@ -22,6 +22,28 @@ use std::process::exit; use nix::unistd::User; +nix::ioctl_read_bad!(tiocgwinsz, libc::TIOCGWINSZ, libc::winsize); + +fn terminal_width() -> Option { + if let Ok(cols) = std::env::var("COLUMNS") { + if let Ok(w) = cols.parse::() { + if w > 0 { + return Some(w); + } + } + } + unsafe { + let mut ws = std::mem::MaybeUninit::::uninit(); + tiocgwinsz(libc::STDOUT_FILENO, ws.as_mut_ptr()).ok()?; + let ws = ws.assume_init(); + if ws.ws_col > 0 { + Some(ws.ws_col as usize) + } else { + None + } + } +} + struct GraphChars { last: &'static str, non_last: &'static str, @@ -110,11 +132,16 @@ fn build_proc_maps() -> Result> { Ok((parent_map, child_map, uid_map)) } +struct PrintOpts<'a> { + graph: Option<&'a GraphChars>, + max_width: Option, +} + fn print_tree( pid: u64, parent_map: &HashMap, child_map: &HashMap>, - graph: Option<&GraphChars>, + opts: &PrintOpts, printed: &mut HashSet, ) -> bool { if !parent_map.contains_key(&pid) { @@ -122,25 +149,17 @@ fn print_tree( return false; } let mut cont = Vec::new(); - let indent_level = print_parents(parent_map, pid, graph, &mut cont, printed); - print_children( - child_map, - pid, - indent_level, - graph, - &mut cont, - true, - printed, - ); + let indent_level = print_parents(parent_map, pid, opts, &mut cont, printed); + print_children(child_map, pid, indent_level, opts, &mut cont, true, printed); true } -fn print_all_trees(child_map: &HashMap>, graph: Option<&GraphChars>) { +fn print_all_trees(child_map: &HashMap>, opts: &PrintOpts) { let mut printed = HashSet::new(); if let Some(root_children) = child_map.get(&0) { for pid in root_children { let mut cont = Vec::new(); - print_children(child_map, *pid, 0, graph, &mut cont, true, &mut printed); + print_children(child_map, *pid, 0, opts, &mut cont, true, &mut printed); } } } @@ -149,7 +168,7 @@ fn print_all_trees(child_map: &HashMap>, graph: Option<&GraphChars fn print_parents( parent_map: &HashMap, pid: u64, - graph: Option<&GraphChars>, + opts: &PrintOpts, cont: &mut Vec, printed: &mut HashSet, ) -> u64 { @@ -166,8 +185,8 @@ fn print_parents( return 0; } - let indent_level = print_parents(parent_map, ppid, graph, cont, printed); - print_ptree_line(ppid, indent_level, graph, cont, true); + let indent_level = print_parents(parent_map, ppid, opts, cont, printed); + print_ptree_line(ppid, indent_level, opts, cont, true); printed.insert(ppid); // Ancestor path is a single chain, no siblings to continue cont.push(false); @@ -178,12 +197,12 @@ fn print_children( child_map: &HashMap>, pid: u64, indent_level: u64, - graph: Option<&GraphChars>, + opts: &PrintOpts, cont: &mut Vec, is_last: bool, printed: &mut HashSet, ) { - print_ptree_line(pid, indent_level, graph, cont, is_last); + print_ptree_line(pid, indent_level, opts, cont, is_last); printed.insert(pid); if let Some(children) = child_map.get(&pid) { for (i, child) in children.iter().enumerate() { @@ -193,7 +212,7 @@ fn print_children( child_map, *child, indent_level + 1, - graph, + opts, cont, child_is_last, printed, @@ -203,37 +222,40 @@ fn print_children( } } -fn print_ptree_line( - pid: u64, - indent_level: u64, - graph: Option<&GraphChars>, - cont: &[bool], - is_last: bool, -) { - match graph { +fn print_ptree_line(pid: u64, indent_level: u64, opts: &PrintOpts, cont: &[bool], is_last: bool) { + use std::fmt::Write; + + let mut line = String::new(); + match opts.graph { Some(g) if indent_level > 0 => { for c in cont.iter().take(indent_level as usize - 1) { if *c { - print!("{}", g.pipe); + line.push_str(g.pipe); } else { - print!("{}", g.space); + line.push_str(g.space); } } if is_last { - print!("{}", g.last); + line.push_str(g.last); } else { - print!("{}", g.non_last); + line.push_str(g.non_last); } } _ => { for _ in 0..indent_level { - print!(" "); + line.push_str(" "); } } } - print!("{pid} "); + let _ = write!(line, "{pid} "); let handle = ptools::proc::ProcHandle::from_pid(pid); - ptools::display::print_cmd_summary_from(&handle); + line.push_str(&ptools::display::cmd_summary_from(&handle)); + if let Some(w) = opts.max_width { + if let Some((idx, _)) = line.char_indices().nth(w) { + line.truncate(idx); + } + } + println!("{line}"); } fn pids_for_user(username: &str, uid_map: &HashMap) -> Result, Box> { @@ -257,11 +279,12 @@ fn pids_for_user(username: &str, uid_map: &HashMap) -> Result struct Args { all: bool, graph: bool, + wrap: bool, target: Vec, } fn print_usage() { - eprintln!("Usage: ptree [-ag] [pid|user]..."); + eprintln!("Usage: ptree [-agw] [pid|user]..."); eprintln!("Print process trees. A /proc/pid path may be used in place of a PID."); eprintln!(); eprintln!("Options:"); @@ -269,6 +292,7 @@ fn print_usage() { eprintln!(" -g, --graph Use line drawing characters"); eprintln!(" -h, --help Print help"); eprintln!(" -V, --version Print version"); + eprintln!(" -w, --wrap Allow output lines to wrap"); } fn parse_args() -> Args { @@ -277,6 +301,7 @@ fn parse_args() -> Args { let mut args = Args { all: false, graph: false, + wrap: false, target: Vec::new(), }; let mut parser = lexopt::Parser::from_env(); @@ -296,6 +321,7 @@ fn parse_args() -> Args { } Short('a') | Long("all") => args.all = true, Short('g') | Long("graph") => args.graph = true, + Short('w') | Long("wrap") => args.wrap = true, Value(val) => { args.target.push(val.to_string_lossy().into_owned()); } @@ -321,10 +347,13 @@ fn main() { } }; - let graph = if args.graph { - Some(graph_chars()) - } else { - None + let opts = PrintOpts { + graph: if args.graph { + Some(graph_chars()) + } else { + None + }, + max_width: if args.wrap { None } else { terminal_width() }, }; let mut error = false; @@ -338,7 +367,7 @@ fn main() { continue; } if printed.insert(pid) - && !print_tree(pid, &parent_map, &child_map, graph, &mut printed) + && !print_tree(pid, &parent_map, &child_map, &opts, &mut printed) { error = true; } @@ -346,7 +375,7 @@ fn main() { } if let Some(pid) = ptools::proc::parse_proc_path(target) { if printed.insert(pid) - && !print_tree(pid, &parent_map, &child_map, graph, &mut printed) + && !print_tree(pid, &parent_map, &child_map, &opts, &mut printed) { error = true; } @@ -360,7 +389,7 @@ fn main() { Ok(pids) => { for pid in pids { if printed.insert(pid) - && !print_tree(pid, &parent_map, &child_map, graph, &mut printed) + && !print_tree(pid, &parent_map, &child_map, &opts, &mut printed) { error = true; } @@ -373,12 +402,12 @@ fn main() { } } } else if args.all { - print_all_trees(&child_map, graph); + print_all_trees(&child_map, &opts); } else if let Some(children) = child_map.get(&1) { let mut printed = HashSet::new(); for pid in children { let mut cont = Vec::new(); - print_children(&child_map, *pid, 0, graph, &mut cont, true, &mut printed); + print_children(&child_map, *pid, 0, &opts, &mut cont, true, &mut printed); } } if error { diff --git a/src/display.rs b/src/display.rs index 3dbce43..8185161 100644 --- a/src/display.rs +++ b/src/display.rs @@ -165,39 +165,35 @@ pub fn print_auxv_from(handle: &ProcHandle) -> io::Result<()> { pub fn print_proc_summary_from(handle: &ProcHandle) { print!("{}:\t", handle.pid()); - print_cmd_summary_from(handle); + println!("{}", cmd_summary_from(handle)); } -pub fn print_cmd_summary_from(handle: &ProcHandle) { +pub fn cmd_summary_from(handle: &ProcHandle) -> String { match handle.read_cmdline() { Ok((args, _)) if !args.is_empty() => { let summary: Vec<_> = args.iter().map(|a| a.to_string_lossy()).collect(); - println!("{}", summary.join(" ")); + summary.join(" ") } Ok(_) => { // Empty cmdline -- fall back to comm name. let is_zombie = matches!(handle.state(), Ok(crate::model::stat::ProcState::Zombie)); match handle.comm() { Ok(ref comm) if !comm.is_empty() => { - print!("{}", comm.to_string_lossy()); + let mut s = comm.to_string_lossy().into_owned(); if is_zombie { - print!(" "); + s.push_str(" "); } - println!(); + s } - Ok(_) => println!(""), - Err(ref e) if e.kind() == io::ErrorKind::NotFound => { - println!(""); - } - Err(_) => println!(""), + Ok(_) => "".to_string(), + Err(ref e) if e.kind() == io::ErrorKind::NotFound => "".to_string(), + Err(_) => "".to_string(), } } - Err(ref e) if e.kind() == io::ErrorKind::NotFound => { - println!(""); - } + Err(ref e) if e.kind() == io::ErrorKind::NotFound => "".to_string(), Err(e) => { - println!(""); eprintln!("{e}"); + "".to_string() } } } diff --git a/tests/ptree_test.rs b/tests/ptree_test.rs index e48749d..36d168b 100644 --- a/tests/ptree_test.rs +++ b/tests/ptree_test.rs @@ -211,6 +211,96 @@ fn ptree_accepts_username_operand() { ); } +#[test] +fn ptree_truncates_lines_to_columns_width() { + let columns = "40"; + let parent_arg = "TRUNCATION_TEST_PARENT_LONG_ARGUMENT"; + let child_arg = "TRUNCATION_TEST_CHILD_LONG_ARGUMENT"; + + let ready = common::ReadySignal::new(true); + let mut example_cmd = Command::new(common::find_exec("examples/ptree_parent_child")); + example_cmd + .args([parent_arg, child_arg]) + .stdin(Stdio::null()) + .stderr(Stdio::inherit()) + .stdout(Stdio::inherit()); + ready.apply_to_command(&mut example_cmd); + + let mut child = example_cmd.spawn().expect("failed to spawn example"); + ready.wait_for_readiness(&mut child); + + let output = Command::new(common::find_exec("ptree")) + .arg(child.id().to_string()) + .env("COLUMNS", columns) + .stdin(Stdio::null()) + .output() + .expect("failed to run ptree"); + + let _ = child.kill(); + let _ = child.wait(); + ready.cleanup(); + + assert!(output.status.success()); + let stdout = String::from_utf8_lossy(&output.stdout); + let max_width: usize = columns.parse().unwrap(); + for line in stdout.lines() { + assert!( + line.len() <= max_width, + "Line exceeds {max_width} columns ({} chars): {line:?}", + line.len() + ); + } + // Verify the output actually has content (not empty) + assert!( + stdout.lines().count() >= 2, + "Expected at least 2 lines of output:\n{stdout}" + ); +} + +#[test] +fn ptree_wrap_flag_disables_truncation() { + let columns = "40"; + let parent_arg = "WRAP_TEST_PARENT_WITH_A_VERY_LONG_ARGUMENT_THAT_EXCEEDS_FORTY_COLUMNS"; + let child_arg = "WRAP_TEST_CHILD"; + + let ready = common::ReadySignal::new(true); + let mut example_cmd = Command::new(common::find_exec("examples/ptree_parent_child")); + example_cmd + .args([parent_arg, child_arg]) + .stdin(Stdio::null()) + .stderr(Stdio::inherit()) + .stdout(Stdio::inherit()); + ready.apply_to_command(&mut example_cmd); + + let mut child = example_cmd.spawn().expect("failed to spawn example"); + ready.wait_for_readiness(&mut child); + + let output = Command::new(common::find_exec("ptree")) + .args(["-w", &child.id().to_string()]) + .env("COLUMNS", columns) + .stdin(Stdio::null()) + .output() + .expect("failed to run ptree"); + + let _ = child.kill(); + let _ = child.wait(); + ready.cleanup(); + + assert!(output.status.success()); + let stdout = String::from_utf8_lossy(&output.stdout); + let max_width: usize = columns.parse().unwrap(); + let has_long_line = stdout.lines().any(|line| line.len() > max_width); + assert!( + has_long_line, + "Expected at least one line to exceed {max_width} columns with -w flag:\n{stdout}" + ); + assert_contains( + &stdout, + parent_arg, + "Expected full parent argument to appear with -w flag", + ); +} + #[test] fn ptree_accepts_mixed_pid_and_user_operands() { let Some(username) = current_username() else {