diff --git a/src/app/sub.rs b/src/app/sub.rs index 4ee419a6..8648fdcc 100644 --- a/src/app/sub.rs +++ b/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, 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 = 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 = 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, path_stack: &mut Vec<&'a OsStr>, - map: &mut HashMap, Color>, + 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); + 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); diff --git a/src/visualizer.rs b/src/visualizer.rs index 2498d5b3..608d3885 100644 --- a/src/visualizer.rs +++ b/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; diff --git a/src/visualizer/coloring.rs b/src/visualizer/coloring.rs index a5c2a392..84674e97 100644 --- a/src/visualizer/coloring.rs +++ b/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, Color>, + map: HashMap, 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, Color>) -> Self { + pub fn new(ls_colors: LsColors, map: HashMap, 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, 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 = ancestors + .chain(std::iter::once(name)) + .map(OsString::from) + .collect(); let color = if has_children { Some(Color::Directory) } else { diff --git a/tests/usual_cli.rs b/tests/usual_cli.rs index 9f0c3e06..f3ca4093 100644 --- a/tests/usual_cli.rs +++ b/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::::as_ref) - .collect::>(), + path.split('/').map(OsString::from).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); +} + +#[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 = 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::>(), color, ) }));