fix(color): multiple arguments

src/app/sub.rs

@@ -17,7 +17,7 @@ use pipe_trait::Pipe;
 use serde::Serialize;
 use std::{
     collections::HashMap,
-    ffi::OsStr,
+    ffi::{OsStr, OsString},
     io::stdout,
     iter::once,
     path::{Path, PathBuf},
@@ -133,7 +133,7 @@ where
         }
 
         let min_ratio: f32 = min_ratio.into();
-        let (data_tree, deduplication_record) = {
+        let (mut data_tree, deduplication_record) = {
             let mut data_tree = data_tree;
             if min_ratio > 0.0 {
                 data_tree.par_cull_insignificant_data(min_ratio);
@@ -144,11 +144,41 @@ where
             let deduplication_record = hardlinks_handler.deduplicate(&mut data_tree);
             if !only_one_arg {
                 assert_eq!(data_tree.name().as_os_str().to_str(), Some(""));
-                *data_tree.name_mut() = OsStringDisplay::os_string_from("(total)");
             }
             (data_tree, deduplication_record)
         };
 
+        // Build the coloring map while the multi-arg root is still "" (a valid path prefix)
+        // so that file_color receives real filesystem paths.
+        let mut leaf_color_map: Option<HashMap<Vec<OsString>, Color>> = color.as_ref().map(|_| {
+            let mut map = HashMap::new();
+            build_coloring_map(&data_tree, &mut Vec::new(), &mut map);
+            map
+        });
+
+        // Rename the synthetic root and rekey the coloring map to match.
+        if !only_one_arg {
+            *data_tree.name_mut() = OsStringDisplay::os_string_from("(total)");
+            if let Some(map) = &mut leaf_color_map {
+                let total = OsString::from("(total)");
+                let empty = OsString::from("");
+                *map = map
+                    .drain()
+                    .map(|(mut key, color)| {
+                        if key.first() == Some(&empty) {
+                            key[0] = total.clone();
+                        }
+                        (key, color)
+                    })
+                    .collect();
+            }
+        }
+
+        let coloring: Option<Coloring> = color.map(|ls_colors| {
+            let map = leaf_color_map.take().unwrap();
+            Coloring::new(ls_colors, map)
+        });
+
         GLOBAL_STATUS_BOARD.clear_line(0);
 
         if let Some(json_output) = json_output {
@@ -198,12 +228,6 @@ where
                 .or(deduplication_result);
         }
 
-        let coloring: Option<Coloring> = 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,
@@ -291,12 +315,12 @@ where
 fn build_coloring_map<'a>(
     node: &'a DataTree<OsStringDisplay, impl size::Size>,
     path_stack: &mut Vec<&'a OsStr>,
-    map: &mut HashMap<Vec<&'a OsStr>, Color>,
+    map: &mut HashMap<Vec<OsString>, Color>,
 ) {
     path_stack.push(node.name().as_os_str());
     if node.children().is_empty() {
         let color = file_color(&path_stack.iter().collect::<PathBuf>());
-        map.insert(path_stack.clone(), color);
+        map.insert(path_stack.iter().map(|s| s.to_os_string()).collect(), color);
     } else {
         for child in node.children() {
             build_coloring_map(child, path_stack, map);

src/visualizer.rs

@@ -62,7 +62,7 @@ where
     /// 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>>,
+    pub coloring: Option<&'a Coloring>,
 }
 
 mod copy;

src/visualizer/coloring.rs

@@ -1,19 +1,23 @@
 use super::{ChildPosition, TreeHorizontalSlice};
 use crate::ls_colors::LsColors;
 use derive_more::Display;
-use std::{collections::HashMap, ffi::OsStr, fmt};
+use std::{
+    collections::HashMap,
+    ffi::{OsStr, OsString},
+    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> {
+pub struct Coloring {
     ls_colors: LsColors,
-    map: HashMap<Vec<&'a OsStr>, Color>,
+    map: HashMap<Vec<OsString>, Color>,
 }
 
-impl<'a> Coloring<'a> {
+impl Coloring {
     /// Create a new [`Coloring`] from LS_COLORS prefixes and a path-components-to-color map.
-    pub fn new(ls_colors: LsColors, map: HashMap<Vec<&'a OsStr>, Color>) -> Self {
+    pub fn new(ls_colors: LsColors, map: HashMap<Vec<OsString>, Color>) -> Self {
         Coloring { ls_colors, map }
     }
 }
@@ -91,7 +95,7 @@ impl Width for ColoredTreeHorizontalSlice<'_> {
 /// 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>>,
+    coloring: Option<&'b Coloring>,
     ancestors: impl Iterator<Item = &'a OsStr>,
     name: &'a OsStr,
     has_children: bool,
@@ -101,7 +105,10 @@ pub(super) fn maybe_colored_slice<'a, 'b>(
         Some(coloring) => coloring,
         None => return MaybeColoredTreeHorizontalSlice::Colorless(slice),
     };
-    let path_components: Vec<&OsStr> = ancestors.chain(std::iter::once(name)).collect();
+    let path_components: Vec<OsString> = ancestors
+        .chain(std::iter::once(name))
+        .map(OsString::from)
+        .collect();
     let color = if has_children {
         Some(Color::Directory)
     } else {

tests/usual_cli.rs

@@ -28,7 +28,7 @@ use parallel_disk_usage::{
     visualizer::{Color, Coloring},
 };
 #[cfg(unix)]
-use std::{collections::HashMap, ffi::OsStr};
+use std::{collections::HashMap, ffi::OsString};
 
 fn stdio(command: Command) -> Command {
     command
@@ -857,9 +857,98 @@ fn color_always() {
     ];
     let leaf_colors = HashMap::from(leaf_colors.map(|(path, color)| {
         (
-            path.split('/')
-                .map(AsRef::<OsStr>::as_ref)
-                .collect::<Vec<_>>(),
+            path.split('/').map(OsString::from).collect::<Vec<_>>(),
+            color,
+        )
+    }));
+    let coloring = Coloring::new(ls_colors, leaf_colors);
+
+    let visualizer = Visualizer::<OsStringDisplay, _> {
+        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);
+}
+
+#[cfg(unix)]
+#[test]
+fn color_always_multiple_args() {
+    let workspace = SampleWorkspace::simple_tree_with_diverse_kinds();
+
+    let args = [
+        "dir-a",
+        "dir-b",
+        "file-root.txt",
+        "link-dir",
+        "link-file.txt",
+        "empty-dir-1",
+        "empty-dir-2",
+    ];
+
+    let actual = {
+        let mut cmd = Command::new(PDU);
+        cmd = cmd
+            .with_current_dir(&workspace)
+            .with_arg("--color=always")
+            .with_arg("--total-width=100")
+            .with_arg("--min-ratio=0")
+            .with_env("LS_COLORS", LS_COLORS);
+        for arg in &args {
+            cmd = cmd.with_arg(arg);
+        }
+        cmd.pipe(stdio)
+            .output()
+            .expect("spawn command with --color=always and multiple args")
+            .pipe(stdout_text)
+    };
+    eprintln!("ACTUAL:\n{actual}\n");
+
+    let data_tree = args
+        .iter()
+        .map(|name| {
+            let builder = FsTreeBuilder {
+                root: workspace.to_path_buf().join(name),
+                size_getter: DEFAULT_GET_SIZE,
+                hardlinks_recorder: &HardlinkIgnorant,
+                reporter: &ErrorOnlyReporter::new(ErrorReport::SILENT),
+                max_depth: 10,
+            };
+            let mut data_tree: DataTree<OsStringDisplay, _> = 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());
+
+    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)/file-root.txt", Color::Normal),
+        ("(total)/link-dir", Color::Symlink),
+        ("(total)/link-file.txt", Color::Symlink),
+        ("(total)/empty-dir-1", Color::Directory),
+        ("(total)/empty-dir-2", Color::Directory),
+    ];
+    let leaf_colors = HashMap::from(leaf_colors.map(|(path, color)| {
+        (
+            path.split('/').map(OsString::from).collect::<Vec<_>>(),
             color,
         )
     }));