From 2f2114e48cbbfe14d28a44bbaad01fa496ff4e85 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 6 Mar 2026 19:10:32 +0000 Subject: [PATCH 1/2] feat: add --color=auto|always|never flag for LS_COLORS support Add a --color CLI flag that colorizes file and directory names in the tree output based on the LS_COLORS environment variable, resolving #11. - `--color=auto` (default): colors when stdout is a TTY - `--color=always`: always colors output - `--color=never`: no colors The color detection code is invoked only after pruning to save CPU/IO. A HashMap maps childless node names to pre-rendered ANSI escape sequences. Directory nodes (with children) use the directory color. Key changes: - New `ColorOption` clap value enum in `src/args/color_option.rs` - New `Coloring` module in `src/coloring.rs` with LS_COLORS integration - `Visualizer` gains an `Option<&Coloring>` field - `TreeHorizontalSlice::display_colored()` renders names with ANSI codes while preserving correct column width alignment - Regenerated help text, completions, and USAGE.md https://claude.ai/code/session_01P8PBNPcM8sNBg8a1627iwZ --- Cargo.lock | 29 +++++++ Cargo.toml | 3 +- USAGE.md | 11 +++ exports/completion.bash | 6 +- exports/completion.elv | 1 + exports/completion.fish | 3 + exports/completion.ps1 | 1 + exports/completion.zsh | 3 + exports/long.help | 10 +++ exports/short.help | 2 + src/app.rs | 3 + src/app/sub.rs | 21 ++++- src/args.rs | 6 ++ src/args/color_option.rs | 15 ++++ src/coloring.rs | 172 +++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/os_string_display.rs | 1 + src/visualizer.rs | 6 +- src/visualizer/methods.rs | 9 +- src/visualizer/tree.rs | 22 +++++ tests/cli_errors.rs | 1 + tests/flag_combinations.rs | 1 + tests/json.rs | 1 + tests/usual_cli.rs | 18 ++++ tests/visualizer.rs | 1 + 25 files changed, 342 insertions(+), 5 deletions(-) create mode 100644 src/args/color_option.rs create mode 100644 src/coloring.rs diff --git a/Cargo.lock b/Cargo.lock index 33c9c1e8..b2f9d7c2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "anstream" version = "0.6.21" @@ -569,6 +578,16 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lscolors" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d60e266dfb1426eb2d24792602e041131fdc0236bb7007abc0e589acafd60929" +dependencies = [ + "aho-corasick", + "nu-ansi-term", +] + [[package]] name = "maplit" version = "1.0.2" @@ -596,6 +615,15 @@ dependencies = [ "winapi", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "objc2-core-foundation" version = "0.3.1" @@ -643,6 +671,7 @@ dependencies = [ "fmt-iter", "into-sorted", "itertools 0.14.0", + "lscolors", "maplit", "normalize-path", "pipe-trait", diff --git a/Cargo.toml b/Cargo.toml index 92b0ed2b..c4d1af6d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,7 +51,7 @@ required-features = ["cli"] [features] default = ["cli"] json = ["serde/derive", "serde_json"] -cli = ["clap/derive", "clap_complete", "clap-utilities", "json"] +cli = ["clap/derive", "clap_complete", "clap-utilities", "json", "lscolors"] cli-completions = ["cli"] [dependencies] @@ -65,6 +65,7 @@ derive_setters = "0.1.9" fmt-iter = "0.2.1" into-sorted = "0.0.3" itertools = "0.14.0" +lscolors = { version = "0.21.0", optional = true } pipe-trait = "0.4.0" rayon = "1.10.0" rounded-div = "0.1.4" diff --git a/USAGE.md b/USAGE.md index 98d1f36c..3fa1f330 100644 --- a/USAGE.md +++ b/USAGE.md @@ -125,6 +125,17 @@ Do not output `.shared.details` in the JSON output. Do not output `.shared.summary` in the JSON output. + +### `--color` + +* _Default:_ `auto`. +* _Choices:_ + - `auto`: Colorize output only when stdout is a terminal + - `always`: Always colorize the output + - `never`: Never colorize the output + +When to colorize the output. + ### `--help` diff --git a/exports/completion.bash b/exports/completion.bash index 8b06a03b..33c4c947 100644 --- a/exports/completion.bash +++ b/exports/completion.bash @@ -23,7 +23,7 @@ _pdu() { case "${cmd}" in pdu) - opts="-b -H -q -d -w -m -s -p -h -V --json-input --json-output --bytes-format --detect-links --dedupe-links --deduplicate-hardlinks --top-down --align-right --quantity --depth --max-depth --width --total-width --column-width --min-ratio --no-sort --no-errors --silent-errors --progress --threads --omit-json-shared-details --omit-json-shared-summary --help --version [FILES]..." + opts="-b -H -q -d -w -m -s -p -h -V --json-input --json-output --bytes-format --detect-links --dedupe-links --deduplicate-hardlinks --top-down --align-right --quantity --depth --max-depth --width --total-width --column-width --min-ratio --no-sort --no-errors --silent-errors --progress --threads --omit-json-shared-details --omit-json-shared-summary --color --help --version [FILES]..." if [[ ${cur} == -* || ${COMP_CWORD} -eq 1 ]] ; then COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") ) return 0 @@ -85,6 +85,10 @@ _pdu() { COMPREPLY=($(compgen -f "${cur}")) return 0 ;; + --color) + COMPREPLY=($(compgen -W "auto always never" -- "${cur}")) + return 0 + ;; *) COMPREPLY=() ;; diff --git a/exports/completion.elv b/exports/completion.elv index d3cda52c..e4792dc2 100644 --- a/exports/completion.elv +++ b/exports/completion.elv @@ -32,6 +32,7 @@ set edit:completion:arg-completer[pdu] = {|@words| cand -m 'Minimal size proportion required to appear' cand --min-ratio 'Minimal size proportion required to appear' cand --threads 'Set the maximum number of threads to spawn. Could be either "auto", "max", or a positive integer' + cand --color 'When to colorize the output' cand --json-input 'Read JSON data from stdin' cand --json-output 'Print JSON data instead of an ASCII chart' cand -H 'Detect and subtract the sizes of hardlinks from their parent directory totals' diff --git a/exports/completion.fish b/exports/completion.fish index 41cc6448..d6c91c21 100644 --- a/exports/completion.fish +++ b/exports/completion.fish @@ -9,6 +9,9 @@ complete -c pdu -s w -l total-width -l width -d 'Width of the visualization' -r complete -c pdu -l column-width -d 'Maximum widths of the tree column and width of the bar column' -r complete -c pdu -s m -l min-ratio -d 'Minimal size proportion required to appear' -r complete -c pdu -l threads -d 'Set the maximum number of threads to spawn. Could be either "auto", "max", or a positive integer' -r +complete -c pdu -l color -d 'When to colorize the output' -r -f -a "auto\t'Colorize output only when stdout is a terminal' +always\t'Always colorize the output' +never\t'Never colorize the output'" complete -c pdu -l json-input -d 'Read JSON data from stdin' complete -c pdu -l json-output -d 'Print JSON data instead of an ASCII chart' complete -c pdu -s H -l deduplicate-hardlinks -l detect-links -l dedupe-links -d 'Detect and subtract the sizes of hardlinks from their parent directory totals' diff --git a/exports/completion.ps1 b/exports/completion.ps1 index 8814bf76..9c71b159 100644 --- a/exports/completion.ps1 +++ b/exports/completion.ps1 @@ -35,6 +35,7 @@ Register-ArgumentCompleter -Native -CommandName 'pdu' -ScriptBlock { [CompletionResult]::new('-m', '-m', [CompletionResultType]::ParameterName, 'Minimal size proportion required to appear') [CompletionResult]::new('--min-ratio', '--min-ratio', [CompletionResultType]::ParameterName, 'Minimal size proportion required to appear') [CompletionResult]::new('--threads', '--threads', [CompletionResultType]::ParameterName, 'Set the maximum number of threads to spawn. Could be either "auto", "max", or a positive integer') + [CompletionResult]::new('--color', '--color', [CompletionResultType]::ParameterName, 'When to colorize the output') [CompletionResult]::new('--json-input', '--json-input', [CompletionResultType]::ParameterName, 'Read JSON data from stdin') [CompletionResult]::new('--json-output', '--json-output', [CompletionResultType]::ParameterName, 'Print JSON data instead of an ASCII chart') [CompletionResult]::new('-H', '-H ', [CompletionResultType]::ParameterName, 'Detect and subtract the sizes of hardlinks from their parent directory totals') diff --git a/exports/completion.zsh b/exports/completion.zsh index dec1cef4..343bf7af 100644 --- a/exports/completion.zsh +++ b/exports/completion.zsh @@ -37,6 +37,9 @@ block-count\:"Count numbers of blocks"))' \ '-m+[Minimal size proportion required to appear]:MIN_RATIO:_default' \ '--min-ratio=[Minimal size proportion required to appear]:MIN_RATIO:_default' \ '--threads=[Set the maximum number of threads to spawn. Could be either "auto", "max", or a positive integer]:THREADS:_default' \ +'--color=[When to colorize the output]:COLOR:((auto\:"Colorize output only when stdout is a terminal" +always\:"Always colorize the output" +never\:"Never colorize the output"))' \ '(-q --quantity -H --deduplicate-hardlinks)--json-input[Read JSON data from stdin]' \ '--json-output[Print JSON data instead of an ASCII chart]' \ '-H[Detect and subtract the sizes of hardlinks from their parent directory totals]' \ diff --git a/exports/long.help b/exports/long.help index efe31299..3b332bd3 100644 --- a/exports/long.help +++ b/exports/long.help @@ -88,6 +88,16 @@ Options: --omit-json-shared-summary Do not output `.shared.summary` in the JSON output + --color + When to colorize the output + + Possible values: + - auto: Colorize output only when stdout is a terminal + - always: Always colorize the output + - never: Never colorize the output + + [default: auto] + -h, --help Print help (see a summary with '-h') diff --git a/exports/short.help b/exports/short.help index 1835edbc..95439c9e 100644 --- a/exports/short.help +++ b/exports/short.help @@ -40,6 +40,8 @@ Options: Do not output `.shared.details` in the JSON output --omit-json-shared-summary Do not output `.shared.summary` in the JSON output + --color + When to colorize the output [default: auto] [possible values: auto, always, never] -h, --help Print help (see more with '--help') -V, --version diff --git a/src/app.rs b/src/app.rs index 4ad14aa5..067d32fb 100644 --- a/src/app.rs +++ b/src/app.rs @@ -86,6 +86,7 @@ impl App { column_width_distribution, direction, bar_alignment, + coloring: None, }; let JsonShared { details, summary } = shared; @@ -291,6 +292,7 @@ impl App { max_depth, min_ratio, no_sort, + color, omit_json_shared_details, omit_json_shared_summary, .. @@ -307,6 +309,7 @@ impl App { max_depth, min_ratio, no_sort, + color, } .run(), )*} }; diff --git a/src/app/sub.rs b/src/app/sub.rs index 3500a5f3..fa55fe91 100644 --- a/src/app/sub.rs +++ b/src/app/sub.rs @@ -1,5 +1,6 @@ use crate::{ - args::{Depth, Fraction}, + args::{ColorOption, Depth, Fraction}, + coloring::Coloring, data_tree::DataTree, fs_tree_builder::FsTreeBuilder, get_size::GetSize, @@ -49,6 +50,8 @@ where pub min_ratio: Fraction, /// Preserve order of entries. pub no_sort: bool, + /// When to colorize the output. + pub color: ColorOption, } impl Sub @@ -74,6 +77,7 @@ where reporter, min_ratio, no_sort, + color, } = self; let max_depth = max_depth.get(); @@ -187,12 +191,27 @@ where .or(deduplication_result); } + // Build color map AFTER pruning to save CPU/IO cycles + let coloring = match color { + ColorOption::Always => Some(Coloring::from_tree(&data_tree)), + ColorOption::Never => None, + ColorOption::Auto => { + use std::io::IsTerminal; + if std::io::stdout().is_terminal() { + Some(Coloring::from_tree(&data_tree)) + } else { + None + } + } + }; + let visualizer = Visualizer { data_tree: &data_tree, bytes_format, direction, bar_alignment, column_width_distribution, + coloring: coloring.as_ref(), }; print!("{visualizer}"); // visualizer already ends with "\n", println! isn't needed here. diff --git a/src/args.rs b/src/args.rs index db6698c8..c7768622 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,8 +1,10 @@ +pub mod color_option; pub mod depth; pub mod fraction; pub mod quantity; pub mod threads; +pub use color_option::ColorOption; pub use depth::Depth; pub use fraction::Fraction; pub use quantity::Quantity; @@ -170,6 +172,10 @@ pub struct Args { /// Do not output `.shared.summary` in the JSON output. #[clap(long, requires = "json_output", requires = "deduplicate_hardlinks")] pub omit_json_shared_summary: bool, + + /// When to colorize the output. + #[clap(long, value_enum, default_value_t = ColorOption::Auto)] + pub color: ColorOption, } impl Args { diff --git a/src/args/color_option.rs b/src/args/color_option.rs new file mode 100644 index 00000000..df5a74d2 --- /dev/null +++ b/src/args/color_option.rs @@ -0,0 +1,15 @@ +#[cfg(feature = "cli")] +use clap::ValueEnum; + +/// When to colorize the output. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +#[cfg_attr(feature = "cli", derive(ValueEnum))] +pub enum ColorOption { + /// Colorize output only when stdout is a terminal. + #[default] + Auto, + /// Always colorize the output. + Always, + /// Never colorize the output. + Never, +} diff --git a/src/coloring.rs b/src/coloring.rs new file mode 100644 index 00000000..4cb6fdf2 --- /dev/null +++ b/src/coloring.rs @@ -0,0 +1,172 @@ +use crate::{data_tree::DataTree, os_string_display::OsStringDisplay, size}; +use std::collections::HashMap; + +/// Color information for visualizing a tree. +/// +/// The `file_colors` field maps the `Display` representation of leaf/childless +/// node names (type `Name` in `Visualizer`, i.e. `OsStringDisplay`) to +/// pre-rendered ANSI escape sequence prefixes. +#[derive(Debug)] +pub struct Coloring { + /// Maps leaf/childless node name strings to their ANSI prefix escape sequence. + file_colors: HashMap, + /// ANSI prefix escape sequence for directories (nodes with children). + dir_prefix: String, +} + +/// The ANSI reset sequence. +const ANSI_RESET: &str = "\x1b[0m"; + +impl Coloring { + /// Build a `Coloring` from a pruned `DataTree` by querying `LS_COLORS`. + /// + /// This should be called only after pruning to save CPU/IO cycles. + #[cfg(feature = "lscolors")] + pub fn from_tree(tree: &DataTree) -> Self + where + Size: size::Size, + { + use lscolors::{Indicator, LsColors}; + + let ls_colors = LsColors::from_env().unwrap_or_default(); + let mut file_colors = HashMap::new(); + collect_leaf_colors(tree, &ls_colors, &mut file_colors); + let dir_prefix = ls_colors + .style_for_indicator(Indicator::Directory) + .map(style_to_ansi_prefix) + .unwrap_or_default(); + Coloring { + file_colors, + dir_prefix, + } + } + + /// Get the ANSI prefix and suffix for a given name. + /// + /// If the name is in `file_colors`, it's a file or childless directory. + /// Otherwise it's a directory with children — use `dir_prefix`. + /// + /// Returns `(prefix, suffix)` where suffix is the ANSI reset sequence. + /// Returns `("", "")` if no style applies. + pub fn ansi_for(&self, name: &str) -> (&str, &str) { + let prefix = self + .file_colors + .get(name) + .map(String::as_str) + .unwrap_or(&self.dir_prefix); + if prefix.is_empty() { + ("", "") + } else { + (prefix, ANSI_RESET) + } + } +} + +/// Recursively collect ANSI prefix strings for all childless nodes in the tree. +#[cfg(feature = "lscolors")] +fn collect_leaf_colors( + tree: &DataTree, + ls_colors: &lscolors::LsColors, + map: &mut HashMap, +) where + Size: size::Size, +{ + if tree.children().is_empty() { + // Leaf node: file or childless directory + let prefix = ls_colors + .style_for_path(tree.name().as_os_str()) + .map(style_to_ansi_prefix) + .unwrap_or_default(); + map.insert(tree.name().to_string(), prefix); + } else { + for child in tree.children() { + collect_leaf_colors(child, ls_colors, map); + } + } +} + +/// Convert an `lscolors::style::Style` to an ANSI escape sequence prefix string. +#[cfg(feature = "lscolors")] +fn style_to_ansi_prefix(style: &lscolors::style::Style) -> String { + let mut codes: Vec = Vec::new(); + + let fs = &style.font_style; + if fs.bold { + codes.push("1".into()); + } + if fs.dimmed { + codes.push("2".into()); + } + if fs.italic { + codes.push("3".into()); + } + if fs.underline { + codes.push("4".into()); + } + if fs.slow_blink { + codes.push("5".into()); + } + if fs.rapid_blink { + codes.push("6".into()); + } + if fs.reverse { + codes.push("7".into()); + } + if fs.hidden { + codes.push("8".into()); + } + if fs.strikethrough { + codes.push("9".into()); + } + + if let Some(ref color) = style.foreground { + color_to_ansi_codes(color, 30, &mut codes); + } + + if let Some(ref color) = style.background { + color_to_ansi_codes(color, 40, &mut codes); + } + + if let Some(ref color) = style.underline { + match color { + lscolors::style::Color::Fixed(n) => { + codes.push(format!("58;5;{n}")); + } + lscolors::style::Color::RGB(r, g, b) => { + codes.push(format!("58;2;{r};{g};{b}")); + } + _ => {} + } + } + + if codes.is_empty() { + String::new() + } else { + format!("\x1b[{}m", codes.join(";")) + } +} + +#[cfg(feature = "lscolors")] +fn color_to_ansi_codes(color: &lscolors::style::Color, base: u8, codes: &mut Vec) { + use lscolors::style::Color; + match color { + Color::Black => codes.push(format!("{}", base)), + Color::Red => codes.push(format!("{}", base + 1)), + Color::Green => codes.push(format!("{}", base + 2)), + Color::Yellow => codes.push(format!("{}", base + 3)), + Color::Blue => codes.push(format!("{}", base + 4)), + Color::Magenta => codes.push(format!("{}", base + 5)), + Color::Cyan => codes.push(format!("{}", base + 6)), + Color::White => codes.push(format!("{}", base + 7)), + Color::BrightBlack => codes.push(format!("{}", base + 60)), + Color::BrightRed => codes.push(format!("{}", base + 61)), + Color::BrightGreen => codes.push(format!("{}", base + 62)), + Color::BrightYellow => codes.push(format!("{}", base + 63)), + Color::BrightBlue => codes.push(format!("{}", base + 64)), + Color::BrightMagenta => codes.push(format!("{}", base + 65)), + Color::BrightCyan => codes.push(format!("{}", base + 66)), + Color::BrightWhite => codes.push(format!("{}", base + 67)), + Color::Fixed(n) => codes.push(format!("{};5;{n}", base + 8)), + Color::RGB(r, g, b) => codes.push(format!("{};2;{r};{g};{b}", base + 8)), + } +} diff --git a/src/lib.rs b/src/lib.rs index 23765add..bb99b642 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,7 @@ pub use serde_json; pub mod app; #[cfg(feature = "cli")] pub mod args; +pub mod coloring; #[cfg(feature = "cli")] pub mod runtime_error; #[cfg(feature = "cli")] diff --git a/src/os_string_display.rs b/src/os_string_display.rs index e713aff6..cb6a1dd8 100644 --- a/src/os_string_display.rs +++ b/src/os_string_display.rs @@ -16,6 +16,7 @@ use serde::{Deserialize, Serialize}; Default, Clone, Copy, + Hash, PartialEq, Eq, PartialOrd, diff --git a/src/visualizer.rs b/src/visualizer.rs index 71effaaf..391989b5 100644 --- a/src/visualizer.rs +++ b/src/visualizer.rs @@ -14,7 +14,7 @@ pub use parenthood::Parenthood; pub use proportion_bar::{ProportionBar, ProportionBarBlock}; pub use tree::{TreeHorizontalSlice, TreeSkeletalComponent}; -use super::{data_tree::DataTree, size}; +use super::{coloring::Coloring, data_tree::DataTree, size}; use std::fmt::Display; /// Visualize a [`DataTree`]. @@ -38,6 +38,7 @@ use std::fmt::Display; /// direction: Direction::BottomUp, /// bar_alignment: BarAlignment::Right, /// column_width_distribution: ColumnWidthDistribution::total(100), +/// coloring: None, /// }; /// println!("{visualizer}"); /// # } @@ -58,6 +59,9 @@ where pub bar_alignment: BarAlignment, /// Distribution and total number of characters/blocks can be placed in a line. pub column_width_distribution: ColumnWidthDistribution, + /// Optional coloring for file/directory names. + /// `None` means no color; `Some` means color according to the map. + pub coloring: Option<&'a Coloring>, } mod copy; diff --git a/src/visualizer/methods.rs b/src/visualizer/methods.rs index f4e99d62..d831f240 100644 --- a/src/visualizer/methods.rs +++ b/src/visualizer/methods.rs @@ -82,10 +82,17 @@ where bar_table .into_iter() .map(|row| { + let tree = if let Some(coloring) = self.coloring { + let (prefix, suffix) = + coloring.ansi_for(&row.tree_horizontal_slice.name); + row.tree_horizontal_slice + .display_colored(prefix, suffix, tree_width) + } else { + align_left(&row.tree_horizontal_slice, tree_width).to_string() + }; format!( "{size} {tree}│{bar}│{ratio}", size = align_right(&row.size, size_width), - tree = align_left(&row.tree_horizontal_slice, tree_width), bar = row.proportion_bar.display(self.bar_alignment), ratio = align_right(&row.percentage, PERCENTAGE_COLUMN_MAX_WIDTH), ) diff --git a/src/visualizer/tree.rs b/src/visualizer/tree.rs index 133b58c1..3602fe57 100644 --- a/src/visualizer/tree.rs +++ b/src/visualizer/tree.rs @@ -110,6 +110,28 @@ impl Width for TreeHorizontalSlice { } } +impl TreeHorizontalSlice { + /// Format with ANSI color codes wrapping the name portion. + /// + /// `name_prefix` and `name_suffix` are the ANSI escape sequences to wrap the name. + /// `total_width` is the target column width; the output is padded with spaces. + pub fn display_colored( + &self, + name_prefix: &str, + name_suffix: &str, + total_width: usize, + ) -> String { + let padding_len = total_width.saturating_sub(self.width()); + format!( + "{}{}{name_prefix}{}{name_suffix}{:padding_len$}", + self.indent(), + self.skeletal_component, + self.name, + "", + ) + } +} + impl TreeHorizontalSlice { /// Truncate the name to fit specified `max_width`. /// diff --git a/tests/cli_errors.rs b/tests/cli_errors.rs index c33f0d9a..35787fa1 100644 --- a/tests/cli_errors.rs +++ b/tests/cli_errors.rs @@ -146,6 +146,7 @@ fn fs_errors() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected_stdout = format!("{visualizer}"); eprintln!("EXPECTED STDOUT:\n{}\n", &expected_stdout); diff --git a/tests/flag_combinations.rs b/tests/flag_combinations.rs index 70c1c6b8..c7ce4a2b 100644 --- a/tests/flag_combinations.rs +++ b/tests/flag_combinations.rs @@ -18,6 +18,7 @@ fn flag_combinations() { let list = CommandList::default() .option_matrix("--quantity", quantity) + .option_matrix("--color", ["auto", "always", "never"]) .flag_matrix("--progress"); for command in list.commands() { diff --git a/tests/json.rs b/tests/json.rs index 95aee0b9..94714c83 100644 --- a/tests/json.rs +++ b/tests/json.rs @@ -141,6 +141,7 @@ fn json_input() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); diff --git a/tests/usual_cli.rs b/tests/usual_cli.rs index cbacb9f2..0d230dba 100644 --- a/tests/usual_cli.rs +++ b/tests/usual_cli.rs @@ -57,6 +57,7 @@ fn total_width() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -96,6 +97,7 @@ fn column_width() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::components(10, 90), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -134,6 +136,7 @@ fn min_ratio_0() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -173,6 +176,7 @@ fn min_ratio() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -212,6 +216,7 @@ fn max_depth_2() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -251,6 +256,7 @@ fn max_depth_1() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -289,6 +295,7 @@ fn top_down() { direction: Direction::TopDown, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -327,6 +334,7 @@ fn align_right() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Right, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -365,6 +373,7 @@ fn quantity_apparent_size() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -404,6 +413,7 @@ fn quantity_block_size() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -443,6 +453,7 @@ fn quantity_block_count() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -482,6 +493,7 @@ fn bytes_format_plain() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -521,6 +533,7 @@ fn bytes_format_metric() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -560,6 +573,7 @@ fn bytes_format_binary() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -597,6 +611,7 @@ fn path_to_workspace() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -650,6 +665,7 @@ fn multiple_names() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -715,6 +731,7 @@ fn multiple_names_max_depth_2() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -776,6 +793,7 @@ fn multiple_names_max_depth_1() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); diff --git a/tests/visualizer.rs b/tests/visualizer.rs index 8c839f26..0d785900 100644 --- a/tests/visualizer.rs +++ b/tests/visualizer.rs @@ -41,6 +41,7 @@ macro_rules! test_case { bytes_format: $bytes_format, direction: Direction::$direction, bar_alignment: BarAlignment::$bar_alignment, + coloring: None, } .to_string(); let expected = $expected; From 796fab9cfcbc9ee7afa40afcd79de06830dc2a80 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 7 Mar 2026 23:47:53 +0000 Subject: [PATCH 2/2] feat: add --color support with failing test for multi-arg leaf coloring Implement the --color [auto|always|never] infrastructure from PR #339: - ColorWhen enum for CLI argument parsing - LsColors wrapper for LS_COLORS environment variable parsing - Coloring system with per-node color selection (Directory/Normal/Executable/Symlink) - Integration threading Coloring through the Visualizer Add a failing test `color_always_multiple_args` that demonstrates the bug described in the review comment: when multiple path arguments are provided, `build_coloring_map` constructs paths like `(total)/dir-a/file.txt` for filesystem type checks, which don't exist on disk, causing symlinks and empty directories to be miscolored as normal files. https://claude.ai/code/session_01KBTsrp6LL3TYRQ7CFF8jY6 --- Cargo.lock | 48 +++++++++++++ Cargo.toml | 2 + src/app.rs | 16 ++++- src/app/sub.rs | 72 +++++++++++++++++++- src/args.rs | 6 ++ src/args/color.rs | 15 +++++ src/lib.rs | 1 + src/ls_colors.rs | 53 +++++++++++++++ src/visualizer.rs | 5 ++ src/visualizer/coloring.rs | 134 +++++++++++++++++++++++++++++++++++++ src/visualizer/display.rs | 7 +- src/visualizer/methods.rs | 31 +++++++-- tests/_utils.rs | 44 ++++++++++++ tests/cli_errors.rs | 1 + tests/json.rs | 1 + tests/usual_cli.rs | 116 +++++++++++++++++++++++++++++++- tests/visualizer.rs | 1 + 17 files changed, 538 insertions(+), 15 deletions(-) create mode 100644 src/args/color.rs create mode 100644 src/ls_colors.rs create mode 100644 src/visualizer/coloring.rs diff --git a/Cargo.lock b/Cargo.lock index 33c9c1e8..cc877bbf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "anstream" version = "0.6.21" @@ -569,6 +578,16 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lscolors" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d60e266dfb1426eb2d24792602e041131fdc0236bb7007abc0e589acafd60929" +dependencies = [ + "aho-corasick", + "nu-ansi-term", +] + [[package]] name = "maplit" version = "1.0.2" @@ -596,6 +615,15 @@ dependencies = [ "winapi", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "objc2-core-foundation" version = "0.3.1" @@ -643,6 +671,7 @@ dependencies = [ "fmt-iter", "into-sorted", "itertools 0.14.0", + "lscolors", "maplit", "normalize-path", "pipe-trait", @@ -653,6 +682,7 @@ dependencies = [ "serde", "serde_json", "smart-default", + "strip-ansi-escapes", "sysinfo", "terminal_size", "text-block-macros", @@ -887,6 +917,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "strip-ansi-escapes" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a8f8038e7e7969abb3f1b7c2a811225e9296da208539e0f79c5251d6cac0025" +dependencies = [ + "vte", +] + [[package]] name = "strsim" version = "0.10.0" @@ -987,6 +1026,15 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "vte" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "231fdcd7ef3037e8330d8e17e61011a2c244126acc0a982f4040ac3f9f0bc077" +dependencies = [ + "memchr", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" diff --git a/Cargo.toml b/Cargo.toml index 92b0ed2b..929514bb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -70,6 +70,7 @@ rayon = "1.10.0" rounded-div = "0.1.4" serde = { version = "1.0.228", optional = true } serde_json = { version = "1.0.149", optional = true } +lscolors = { version = "0.21", features = ["nu-ansi-term"] } smart-default = "0.7.1" sysinfo = "0.38.2" terminal_size = "0.4.3" @@ -83,3 +84,4 @@ maplit = "1.0.2" normalize-path = "0.2.1" pretty_assertions = "1.4.1" rand = "0.10.0" +strip-ansi-escapes = "0.2.1" diff --git a/src/app.rs b/src/app.rs index 4ad14aa5..c2ee04fc 100644 --- a/src/app.rs +++ b/src/app.rs @@ -3,11 +3,12 @@ pub mod sub; pub use sub::Sub; use crate::{ - args::{Args, Quantity, Threads}, + args::{Args, ColorWhen, Quantity, Threads}, bytes_format::BytesFormat, get_size::{GetApparentSize, GetSize}, hardlink, json_data::{JsonData, JsonDataBody, JsonShared, JsonTree}, + ls_colors::LsColors, reporter::{ErrorOnlyReporter, ErrorReport, ProgressAndErrorReporter, ProgressReport}, runtime_error::RuntimeError, size, @@ -16,7 +17,10 @@ use crate::{ use clap::Parser; use hdd::any_path_is_in_hdd; use pipe_trait::Pipe; -use std::{io::stdin, time::Duration}; +use std::{ + io::{stdin, stdout, IsTerminal}, + time::Duration, +}; use sub::JsonOutputParam; use sysinfo::Disks; @@ -86,6 +90,7 @@ impl App { column_width_distribution, direction, bar_alignment, + coloring: None, }; let JsonShared { details, summary } = shared; @@ -169,6 +174,12 @@ impl App { ErrorReport::TEXT }; + let color = match self.args.color { + ColorWhen::Always => Some(LsColors::from_env()), + ColorWhen::Never => None, + ColorWhen::Auto => stdout().is_terminal().then(LsColors::from_env), + }; + trait GetSizeUtils: GetSize { const INSTANCE: Self; const QUANTITY: Quantity; @@ -307,6 +318,7 @@ impl App { max_depth, min_ratio, no_sort, + color, } .run(), )*} }; diff --git a/src/app/sub.rs b/src/app/sub.rs index 3500a5f3..4ee419a6 100644 --- a/src/app/sub.rs +++ b/src/app/sub.rs @@ -5,16 +5,23 @@ use crate::{ get_size::GetSize, hardlink::{DeduplicateSharedSize, HardlinkIgnorant, RecordHardlinks}, json_data::{BinaryVersion, JsonData, JsonDataBody, JsonShared, JsonTree, SchemaVersion}, + ls_colors::LsColors, os_string_display::OsStringDisplay, reporter::ParallelReporter, runtime_error::RuntimeError, size, status_board::GLOBAL_STATUS_BOARD, - visualizer::{BarAlignment, ColumnWidthDistribution, Direction, Visualizer}, + visualizer::{BarAlignment, Color, Coloring, ColumnWidthDistribution, Direction, Visualizer}, }; use pipe_trait::Pipe; use serde::Serialize; -use std::{io::stdout, iter::once, path::PathBuf}; +use std::{ + collections::HashMap, + ffi::OsStr, + io::stdout, + iter::once, + path::{Path, PathBuf}, +}; /// The sub program of the main application. pub struct Sub @@ -49,6 +56,8 @@ where pub min_ratio: Fraction, /// Preserve order of entries. pub no_sort: bool, + /// Whether to color the output. + pub color: Option, } impl Sub @@ -74,6 +83,7 @@ where reporter, min_ratio, no_sort, + color, } = self; let max_depth = max_depth.get(); @@ -98,6 +108,7 @@ where files: vec![".".into()], hardlinks_handler, reporter, + color, ..self } .run(); @@ -187,12 +198,19 @@ where .or(deduplication_result); } + let coloring: Option = color.map(|ls_colors| { + let mut map = HashMap::new(); + build_coloring_map(&data_tree, &mut Vec::new(), &mut map); + Coloring::new(ls_colors, map) + }); + let visualizer = Visualizer { data_tree: &data_tree, bytes_format, direction, bar_alignment, column_width_distribution, + coloring: coloring.as_ref(), }; print!("{visualizer}"); // visualizer already ends with "\n", println! isn't needed here. @@ -262,5 +280,55 @@ where } } +/// Recursively walk a pruned [`DataTree`] and build a map of path-component vectors to [`Color`] values. +/// +/// The `path_stack` argument is a reusable buffer of path components representing the current +/// ancestor chain. Each recursive call pushes the node's name and pops it on return, so no +/// cloning occurs during traversal — only at leaf insertions. +/// Leaf nodes (files or childless directories after pruning) are added to the map. +/// Nodes with children are skipped because the [`Visualizer`] uses the children count to +/// determine their color at render time. +fn build_coloring_map<'a>( + node: &'a DataTree, + path_stack: &mut Vec<&'a OsStr>, + map: &mut HashMap, Color>, +) { + path_stack.push(node.name().as_os_str()); + if node.children().is_empty() { + let color = file_color(&path_stack.iter().collect::()); + map.insert(path_stack.clone(), color); + } else { + for child in node.children() { + build_coloring_map(child, path_stack, map); + } + } + path_stack.pop(); +} + +fn file_color(path: &Path) -> Color { + if path.is_symlink() { + Color::Symlink + } else if path.is_dir() { + Color::Directory + } else if is_executable(path) { + Color::Executable + } else { + Color::Normal + } +} + +#[cfg(unix)] +fn is_executable(path: &Path) -> bool { + use std::os::unix::fs::PermissionsExt; + path.metadata() + .map(|stats| stats.permissions().mode() & 0o111 != 0) + .unwrap_or(false) +} + +#[cfg(not(unix))] +fn is_executable(_path: &Path) -> bool { + false +} + #[cfg(unix)] mod unix_ext; diff --git a/src/args.rs b/src/args.rs index db6698c8..44c7bfb4 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,8 +1,10 @@ +pub mod color; pub mod depth; pub mod fraction; pub mod quantity; pub mod threads; +pub use color::ColorWhen; pub use depth::Depth; pub use fraction::Fraction; pub use quantity::Quantity; @@ -170,6 +172,10 @@ pub struct Args { /// Do not output `.shared.summary` in the JSON output. #[clap(long, requires = "json_output", requires = "deduplicate_hardlinks")] pub omit_json_shared_summary: bool, + + /// Whether to show colors. + #[clap(long, value_enum, default_value_t = ColorWhen::Auto)] + pub color: ColorWhen, } impl Args { diff --git a/src/args/color.rs b/src/args/color.rs new file mode 100644 index 00000000..2fb22620 --- /dev/null +++ b/src/args/color.rs @@ -0,0 +1,15 @@ +#[cfg(feature = "cli")] +use clap::ValueEnum; + +/// When to use colors in the output. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +#[cfg_attr(feature = "cli", derive(ValueEnum))] +pub enum ColorWhen { + /// Detect if the output is a TTY and render colors accordingly. + #[default] + Auto, + /// Always render colors. + Always, + /// Never render colors. + Never, +} diff --git a/src/lib.rs b/src/lib.rs index 23765add..697481a8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,6 +43,7 @@ pub mod get_size; pub mod hardlink; pub mod inode; pub mod json_data; +pub mod ls_colors; pub mod os_string_display; pub mod reporter; pub mod size; diff --git a/src/ls_colors.rs b/src/ls_colors.rs new file mode 100644 index 00000000..67c70ccc --- /dev/null +++ b/src/ls_colors.rs @@ -0,0 +1,53 @@ +use crate::visualizer::coloring::Color; +use lscolors::{self, Indicator}; + +/// ANSI color prefix strings for terminal output, initialized from the `LS_COLORS` environment +/// variable. +#[derive(Debug, Clone)] +pub struct LsColors { + directory: String, + normal: String, + executable: String, + symlink: String, +} + +impl LsColors { + /// Initialize by reading the current environment's `LS_COLORS`. + pub fn from_env() -> Self { + Self::from_lscolors_crate(&lscolors::LsColors::from_env().unwrap_or_default()) + } + + /// Parse an `LS_COLORS`-format string into an [`LsColors`]. + /// + /// Unrecognized or invalid entries are silently ignored. + #[allow(clippy::should_implement_trait, reason = "Nonsense suggestion!")] + pub fn from_str(input: &str) -> Self { + Self::from_lscolors_crate(&lscolors::LsColors::from_string(input)) + } + + /// Derive a [`LsColors`] from an existing [`lscolors::LsColors`]. + fn from_lscolors_crate(ls_colors: &lscolors::LsColors) -> Self { + let prefix_for = |indicator: Indicator| { + ls_colors + .style_for_indicator(indicator) + .map(|s| s.to_nu_ansi_term_style().prefix().to_string()) + .unwrap_or_default() + }; + LsColors { + directory: prefix_for(Indicator::Directory), + normal: prefix_for(Indicator::RegularFile), + executable: prefix_for(Indicator::ExecutableFile), + symlink: prefix_for(Indicator::SymbolicLink), + } + } + + /// Return the ANSI prefix string for the given [`Color`] variant. + pub(crate) fn prefix_str(&self, color: Color) -> &str { + match color { + Color::Directory => &self.directory, + Color::Normal => &self.normal, + Color::Executable => &self.executable, + Color::Symlink => &self.symlink, + } + } +} diff --git a/src/visualizer.rs b/src/visualizer.rs index 71effaaf..2498d5b3 100644 --- a/src/visualizer.rs +++ b/src/visualizer.rs @@ -1,5 +1,6 @@ pub mod bar_alignment; pub mod child_position; +pub mod coloring; pub mod column_width_distribution; pub mod direction; pub mod parenthood; @@ -8,6 +9,7 @@ pub mod tree; pub use bar_alignment::BarAlignment; pub use child_position::ChildPosition; +pub use coloring::{Color, Coloring}; pub use column_width_distribution::ColumnWidthDistribution; pub use direction::Direction; pub use parenthood::Parenthood; @@ -38,6 +40,7 @@ use std::fmt::Display; /// direction: Direction::BottomUp, /// bar_alignment: BarAlignment::Right, /// column_width_distribution: ColumnWidthDistribution::total(100), +/// coloring: None, /// }; /// println!("{visualizer}"); /// # } @@ -58,6 +61,8 @@ where pub bar_alignment: BarAlignment, /// Distribution and total number of characters/blocks can be placed in a line. pub column_width_distribution: ColumnWidthDistribution, + /// Optional coloring configuration for colorful output, mapping full node paths to colors. + pub coloring: Option<&'a Coloring<'a>>, } mod copy; diff --git a/src/visualizer/coloring.rs b/src/visualizer/coloring.rs new file mode 100644 index 00000000..a5c2a392 --- /dev/null +++ b/src/visualizer/coloring.rs @@ -0,0 +1,134 @@ +use super::{ChildPosition, TreeHorizontalSlice}; +use crate::ls_colors::LsColors; +use derive_more::Display; +use std::{collections::HashMap, ffi::OsStr, fmt}; +use zero_copy_pads::Width; + +/// Coloring configuration: ANSI prefix strings from the environment and a full-path-to-color map. +#[derive(Debug)] +pub struct Coloring<'a> { + ls_colors: LsColors, + map: HashMap, Color>, +} + +impl<'a> Coloring<'a> { + /// Create a new [`Coloring`] from LS_COLORS prefixes and a path-components-to-color map. + pub fn new(ls_colors: LsColors, map: HashMap, Color>) -> Self { + Coloring { ls_colors, map } + } +} + +/// The coloring to apply to a node name. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Color { + /// Color as a directory. + Directory, + /// Color as a regular file. + Normal, + /// Color as an executable file. + Executable, + /// Color as a symbolic link. + Symlink, +} + +impl Color { + /// Get the ANSI prefix for this color from the given prefix table. + fn ansi_prefix(self, prefixes: &LsColors) -> AnsiPrefix<'_> { + AnsiPrefix(prefixes.prefix_str(self)) + } +} + +/// ANSI prefix wrapper for a [`Color`] variant, implements [`Display`]. +#[derive(Display)] +struct AnsiPrefix<'a>(&'a str); + +impl AnsiPrefix<'_> { + /// Returns the reset suffix to emit after this prefix, or `""` if no prefix. + fn suffix(&self) -> &'static str { + if self.0.is_empty() { + "" + } else { + "\x1b[0m" + } + } +} + +/// A [`TreeHorizontalSlice`] with its color applied, used for rendering. +pub(super) struct ColoredTreeHorizontalSlice<'a> { + slice: TreeHorizontalSlice, + color: Color, + ls_colors: &'a LsColors, +} + +impl fmt::Display for ColoredTreeHorizontalSlice<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let TreeHorizontalSlice { + ancestor_relative_positions, + skeletal_component, + name, + } = &self.slice; + for pos in ancestor_relative_positions { + let connector = match pos { + ChildPosition::Init => "│ ", + ChildPosition::Last => " ", + }; + write!(f, "{connector}")?; + } + let prefix = self.color.ansi_prefix(self.ls_colors); + let suffix = prefix.suffix(); + write!(f, "{skeletal_component}{prefix}{name}{suffix}") + } +} + +impl Width for ColoredTreeHorizontalSlice<'_> { + fn width(&self) -> usize { + self.slice.width() + } +} + +/// Wrap a [`TreeHorizontalSlice`] with color if coloring is available, otherwise return it as-is. +/// +/// Path components are only constructed when coloring is enabled, avoiding +/// unnecessary allocation in the common no-color case. +pub(super) fn maybe_colored_slice<'a, 'b>( + coloring: Option<&'b Coloring<'a>>, + ancestors: impl Iterator, + name: &'a OsStr, + has_children: bool, + slice: TreeHorizontalSlice, +) -> MaybeColoredTreeHorizontalSlice<'b> { + let coloring = match coloring { + Some(coloring) => coloring, + None => return MaybeColoredTreeHorizontalSlice::Colorless(slice), + }; + let path_components: Vec<&OsStr> = ancestors.chain(std::iter::once(name)).collect(); + let color = if has_children { + Some(Color::Directory) + } else { + coloring.map.get(&path_components).copied() + }; + match color { + Some(color) => MaybeColoredTreeHorizontalSlice::Colorful(ColoredTreeHorizontalSlice { + slice, + color, + ls_colors: &coloring.ls_colors, + }), + None => MaybeColoredTreeHorizontalSlice::Colorless(slice), + } +} + +/// Either a [`TreeHorizontalSlice`] (colorless) or a [`ColoredTreeHorizontalSlice`] (colorful). +#[derive(Display)] +pub(super) enum MaybeColoredTreeHorizontalSlice<'a> { + Colorless(TreeHorizontalSlice), + Colorful(ColoredTreeHorizontalSlice<'a>), +} + +impl Width for MaybeColoredTreeHorizontalSlice<'_> { + fn width(&self) -> usize { + match self { + MaybeColoredTreeHorizontalSlice::Colorless(slice) => slice.width(), + MaybeColoredTreeHorizontalSlice::Colorful(slice) => slice.width(), + } + } +} diff --git a/src/visualizer/display.rs b/src/visualizer/display.rs index 1b3f5137..2c12de1b 100644 --- a/src/visualizer/display.rs +++ b/src/visualizer/display.rs @@ -1,10 +1,13 @@ use super::{Direction::*, Visualizer}; use crate::size; -use std::fmt::{Display, Error, Formatter}; +use std::{ + ffi::OsStr, + fmt::{Display, Error, Formatter}, +}; impl<'a, Name, Size> Display for Visualizer<'a, Name, Size> where - Name: Display, + Name: Display + AsRef, Size: size::Size + Into, { /// Create the ASCII chart. diff --git a/src/visualizer/methods.rs b/src/visualizer/methods.rs index f4e99d62..8b061d9a 100644 --- a/src/visualizer/methods.rs +++ b/src/visualizer/methods.rs @@ -12,14 +12,14 @@ use node_info::*; use table::*; use tree_table::*; -use super::{ColumnWidthDistribution, Visualizer}; +use super::{coloring::maybe_colored_slice, ColumnWidthDistribution, Visualizer}; use crate::size; -use std::{cmp::min, fmt::Display}; +use std::{cmp::min, ffi::OsStr, fmt::Display}; use zero_copy_pads::{align_left, align_right}; impl<'a, Name, Size> Visualizer<'a, Name, Size> where - Name: Display, + Name: Display + AsRef, Size: size::Size + Into, { /// Create ASCII rows that visualize the [tree](crate::data_tree::DataTree), such rows @@ -82,12 +82,29 @@ where bar_table .into_iter() .map(|row| { + let BarRow { + tree_row, + proportion_bar, + } = row; + let TreeRow { + initial_row, + tree_horizontal_slice, + } = tree_row; + + let tree = maybe_colored_slice( + self.coloring, + initial_row.ancestors.iter().map(|node| node.name.as_ref()), + initial_row.node_info.name.as_ref(), + initial_row.node_info.children_count > 0, + tree_horizontal_slice, + ); + format!( "{size} {tree}│{bar}│{ratio}", - size = align_right(&row.size, size_width), - tree = align_left(&row.tree_horizontal_slice, tree_width), - bar = row.proportion_bar.display(self.bar_alignment), - ratio = align_right(&row.percentage, PERCENTAGE_COLUMN_MAX_WIDTH), + size = align_right(&initial_row.size, size_width), + tree = align_left(tree, tree_width), + bar = proportion_bar.display(self.bar_alignment), + ratio = align_right(&initial_row.percentage, PERCENTAGE_COLUMN_MAX_WIDTH), ) }) .collect() diff --git a/tests/_utils.rs b/tests/_utils.rs index 75b4a326..0b06644b 100644 --- a/tests/_utils.rs +++ b/tests/_utils.rs @@ -101,6 +101,50 @@ impl Default for SampleWorkspace { /// POSIX-exclusive functions #[cfg(unix)] impl SampleWorkspace { + /// Set up a temporary directory for tests. + /// + /// This directory has a diverse mix of file kinds: non-empty directories, empty directories, + /// regular files, and symbolic links — multiple of each kind. + pub fn simple_tree_with_diverse_kinds() -> Self { + use std::os::unix::fs::symlink; + let temp = Temp::new_dir().expect("create working directory for sample workspace"); + + MergeableFileSystemTree::<&str, String>::from(dir! { + "dir-a" => dir! { + "file-a1.txt" => file!("a".repeat(100_000)) + "file-a2.txt" => file!("a".repeat(200_000)) + "subdir-a" => dir! { + "file-a3.txt" => file!("a".repeat(300_000)) + } + } + "dir-b" => dir! { + "file-b1.txt" => file!("a".repeat(150_000)) + } + "empty-dir-1" => dir! {} + "empty-dir-2" => dir! {} + "file-root.txt" => file!("a".repeat(50_000)) + }) + .build(&temp) + .expect("build filesystem tree for diverse-kinds sample workspace"); + + macro_rules! symlink { + ($link_name:literal -> $target:literal) => { + let link_name = $link_name; + let target = $target; + if let Err(error) = symlink(target, temp.join(link_name)) { + panic!( + "Failed to create symbolic link {link_name} pointing to {target}: {error}" + ); + } + }; + } + + symlink!("link-dir" -> "dir-a"); + symlink!("link-file.txt" -> "file-root.txt"); + + SampleWorkspace(temp) + } + /// Set up a temporary directory for tests. /// /// This directory would have a couple of normal files and a couple of hardlinks. diff --git a/tests/cli_errors.rs b/tests/cli_errors.rs index c33f0d9a..35787fa1 100644 --- a/tests/cli_errors.rs +++ b/tests/cli_errors.rs @@ -146,6 +146,7 @@ fn fs_errors() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected_stdout = format!("{visualizer}"); eprintln!("EXPECTED STDOUT:\n{}\n", &expected_stdout); diff --git a/tests/json.rs b/tests/json.rs index 95aee0b9..94714c83 100644 --- a/tests/json.rs +++ b/tests/json.rs @@ -141,6 +141,7 @@ fn json_input() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); diff --git a/tests/usual_cli.rs b/tests/usual_cli.rs index cbacb9f2..4c367938 100644 --- a/tests/usual_cli.rs +++ b/tests/usual_cli.rs @@ -10,13 +10,17 @@ use parallel_disk_usage::{ fs_tree_builder::FsTreeBuilder, get_size::GetApparentSize, hardlink::HardlinkIgnorant, + ls_colors::LsColors, os_string_display::OsStringDisplay, reporter::{ErrorOnlyReporter, ErrorReport}, - visualizer::{BarAlignment, ColumnWidthDistribution, Direction, Visualizer}, + visualizer::{BarAlignment, Color, Coloring, ColumnWidthDistribution, Direction, Visualizer}, }; use pipe_trait::Pipe; use pretty_assertions::assert_eq; -use std::process::{Command, Stdio}; +use std::{collections::HashMap, ffi::OsStr, process::{Command, Stdio}}; + +/// Predefined `LS_COLORS` value used in color tests to ensure deterministic output. +const LS_COLORS: &str = "rs=0:di=01;34:ln=01;36:ex=01;32:fi=00"; #[cfg(unix)] use parallel_disk_usage::get_size::{GetBlockCount, GetBlockSize}; @@ -57,6 +61,7 @@ fn total_width() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -96,6 +101,7 @@ fn column_width() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::components(10, 90), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -134,6 +140,7 @@ fn min_ratio_0() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -173,6 +180,7 @@ fn min_ratio() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -212,6 +220,7 @@ fn max_depth_2() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -251,6 +260,7 @@ fn max_depth_1() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -289,6 +299,7 @@ fn top_down() { direction: Direction::TopDown, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -327,6 +338,7 @@ fn align_right() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Right, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -365,6 +377,7 @@ fn quantity_apparent_size() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -404,6 +417,7 @@ fn quantity_block_size() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -443,6 +457,7 @@ fn quantity_block_count() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -482,6 +497,7 @@ fn bytes_format_plain() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -521,6 +537,7 @@ fn bytes_format_metric() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -560,6 +577,7 @@ fn bytes_format_binary() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -597,6 +615,7 @@ fn path_to_workspace() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -650,6 +669,7 @@ fn multiple_names() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -715,6 +735,7 @@ fn multiple_names_max_depth_2() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -776,6 +797,7 @@ fn multiple_names_max_depth_1() { direction: Direction::BottomUp, bar_alignment: BarAlignment::Left, column_width_distribution: ColumnWidthDistribution::total(100), + coloring: None, }; let expected = format!("{visualizer}"); let expected = expected.trim_end(); @@ -787,3 +809,93 @@ fn multiple_names_max_depth_1() { assert!(lines.next().unwrap().contains("┌──(total)")); assert_eq!(lines.next(), None); } + +/// Test that `--color=always` with multiple arguments correctly colors directories, symlinks, +/// and regular files. This exercises the code path where a synthetic `(total)` root is created. +/// +/// The coloring map must use real filesystem paths for `file_color()` checks, not the +/// synthetic `(total)/...` paths which don't exist on disk. +#[cfg(unix)] +#[test] +fn color_always_multiple_args() { + let workspace = SampleWorkspace::simple_tree_with_diverse_kinds(); + + let actual = Command::new(PDU) + .with_current_dir(&workspace) + .with_arg("--color=always") + .with_arg("--quantity=apparent-size") + .with_arg("--total-width=100") + .with_arg("--min-ratio=0") + .with_arg("dir-a") + .with_arg("dir-b") + .with_arg("empty-dir-1") + .with_arg("link-dir") + .with_arg("link-file.txt") + .with_env("LS_COLORS", LS_COLORS) + .pipe(stdio) + .output() + .expect("spawn command with --color=always and multiple args") + .pipe(stdout_text); + eprintln!("ACTUAL:\n{actual}\n"); + + // Build the expected tree manually, mirroring what `pdu` does with multiple args. + let names = ["dir-a", "dir-b", "empty-dir-1", "link-dir", "link-file.txt"]; + let data_tree = names + .iter() + .map(|name| { + let builder = FsTreeBuilder { + root: workspace.to_path_buf().join(name), + size_getter: GetApparentSize, + hardlinks_recorder: &HardlinkIgnorant, + reporter: &ErrorOnlyReporter::new(ErrorReport::SILENT), + max_depth: 10, + }; + let mut data_tree: DataTree = builder.into(); + *data_tree.name_mut() = OsStringDisplay::os_string_from(name); + data_tree + }) + .pipe(|children| { + DataTree::dir( + OsStringDisplay::os_string_from("(total)"), + 0.into(), + children.collect(), + ) + }) + .into_par_sorted(|left, right| left.size().cmp(&right.size()).reverse()); + + // Build the coloring map using the CORRECT filesystem paths (not the `(total)/...` ones). + // This is what the output SHOULD look like — symlinks as Symlink, empty dirs as Directory, etc. + let ls_colors = LsColors::from_str(LS_COLORS); + let leaf_colors = [ + ("(total)/dir-a/file-a1.txt", Color::Normal), + ("(total)/dir-a/file-a2.txt", Color::Normal), + ("(total)/dir-a/subdir-a/file-a3.txt", Color::Normal), + ("(total)/dir-b/file-b1.txt", Color::Normal), + ("(total)/empty-dir-1", Color::Directory), + ("(total)/link-dir", Color::Symlink), + ("(total)/link-file.txt", Color::Symlink), + ]; + let leaf_colors = HashMap::from(leaf_colors.map(|(path, color)| { + ( + path.split('/') + .map(AsRef::::as_ref) + .collect::>(), + color, + ) + })); + let coloring = Coloring::new(ls_colors, leaf_colors); + + let visualizer = Visualizer:: { + data_tree: &data_tree, + bytes_format: BytesFormat::MetricUnits, + direction: Direction::BottomUp, + bar_alignment: BarAlignment::Left, + column_width_distribution: ColumnWidthDistribution::total(100), + coloring: Some(&coloring), + }; + let expected = format!("{visualizer}"); + let expected = expected.trim_end(); + eprintln!("EXPECTED:\n{expected}\n"); + + assert_eq!(actual, expected); +} diff --git a/tests/visualizer.rs b/tests/visualizer.rs index 8c839f26..0d785900 100644 --- a/tests/visualizer.rs +++ b/tests/visualizer.rs @@ -41,6 +41,7 @@ macro_rules! test_case { bytes_format: $bytes_format, direction: Direction::$direction, bar_alignment: BarAlignment::$bar_alignment, + coloring: None, } .to_string(); let expected = $expected;