feat: add --color support with failing test for multi-arg leaf coloring

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"
@@ -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"

USAGE.md

@@ -125,6 +125,17 @@ Do not output `.shared.details` in the JSON output.
 
 Do not output `.shared.summary` in the JSON output.
 
+<a id="color" name="color"></a>
+### `--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.
+
 <a id="option-h" name="option-h"></a><a id="help" name="help"></a>
 ### `--help`
 

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=()
                     ;;

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'

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'

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')

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]' \

exports/long.help

@@ -88,6 +88,16 @@ Options:
       --omit-json-shared-summary
           Do not output `.shared.summary` in the JSON output
 
+      --color <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')
 

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 <COLOR>
+          When to colorize the output [default: auto] [possible values: auto, always, never]
   -h, --help
           Print help (see more with '--help')
   -V, --version

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(),
             )*} };

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<Size, SizeGetter, HardlinksHandler, Report> Sub<Size, SizeGetter, HardlinksHandler, Report>
@@ -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.

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 {

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,
+}