diff --git a/README.md b/README.md index 489a3e77..6f05fffd 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,7 @@ Arguments: [DIR] Root directory to traverse; defaults to current working directory Options: + -c, --count Include aggregate file count in tree output -d, --disk-usage Print physical or logical file size [default: logical] [possible values: logical, physical] -g, --glob Include or exclude files using glob patterns --iglob Include or exclude files using glob patterns; case insensitive diff --git a/src/render/context/mod.rs b/src/render/context/mod.rs index 7369615b..89aaca7a 100644 --- a/src/render/context/mod.rs +++ b/src/render/context/mod.rs @@ -26,6 +26,10 @@ mod test; #[command(version = "1.6.0")] #[command(about = "erdtree (et) is a multi-threaded filetree visualizer and disk usage analyzer.", long_about = None)] pub struct Context { + /// Include aggregate file count in tree output + #[arg(short, long)] + pub count: bool, + /// Root directory to traverse; defaults to current working directory dir: Option, diff --git a/src/render/styles.rs b/src/render/styles.rs index 1146b7cd..d8064326 100644 --- a/src/render/styles.rs +++ b/src/render/styles.rs @@ -27,7 +27,7 @@ pub static LS_COLORS: OnceCell = OnceCell::new(); /// Runtime evaluated static that contains ANSI-colored box drawing characters used for the /// printing of [super::tree::Tree]'s branches. -pub static THEME: OnceCell = OnceCell::new(); +pub static TREE_THEME: OnceCell = OnceCell::new(); /// Runtime evaluated static that contains ANSI-colored box drawing characters used for the /// printing of [super::tree::Tree]'s branches for descendents of symlinks. @@ -59,13 +59,13 @@ pub fn get_du_theme() -> &'static HashMap<&'static str, Color> { } /// Getter for [THEME]. Panics if not initialized. -pub fn get_theme() -> &'static ThemesMap { - THEME.get().expect("THEME not initialized") +pub fn get_tree_theme() -> &'static ThemesMap { + TREE_THEME.get().expect("TREE_THEME not initialized") } /// Getter for [LINK_THEME]. Panics if not initialized. pub fn get_link_theme() -> &'static ThemesMap { - LINK_THEME.get().expect("THEME not initialized") + LINK_THEME.get().expect("LINK_THEME not initialized") } /// Initializes [LS_COLORS] by reading in the `LS_COLORS` environment variable. If it isn't set, a @@ -84,7 +84,7 @@ fn init_themes() { "vtrt" => format!("{}", Color::Purple.paint(VTRT)) }; - THEME.set(theme).unwrap(); + TREE_THEME.set(theme).unwrap(); let link_theme = hash! { "vt" => format!("{}", Color::Red.paint(VT)), diff --git a/src/render/tree/count.rs b/src/render/tree/count.rs new file mode 100644 index 00000000..15b07385 --- /dev/null +++ b/src/render/tree/count.rs @@ -0,0 +1,90 @@ +use super::Node; +use std::{ + convert::From, + fmt::{self, Display}, +}; + +/// For keeping track of the number of various file-types of [Node]'s chlidren. +#[derive(Default)] +pub struct FileCount { + pub num_dirs: usize, + pub num_files: usize, + pub num_links: usize, +} + +impl FileCount { + /// Update [Self] with information from [Node]. + pub fn update(&mut self, node: &Node) { + if node.is_dir() { + self.num_dirs += 1; + } else if node.is_symlink() { + self.num_links += 1; + } else { + self.num_files += 1; + } + } + + /// Update [Self] with information from [Self]. + pub fn update_from_count( + &mut self, + Self { + num_dirs, + num_files, + num_links, + }: Self, + ) { + self.num_dirs += num_dirs; + self.num_links += num_links; + self.num_files += num_files; + } +} + +impl From> for FileCount { + fn from(data: Vec) -> Self { + let mut agg = Self::default(); + + for datum in data { + agg.update_from_count(datum); + } + + agg + } +} + +impl Display for FileCount { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut components = vec![]; + + if self.num_dirs > 0 { + let output = if self.num_dirs > 1 { + format!("{} {}", self.num_dirs, "directories") + } else { + format!("{} {}", self.num_dirs, "directory") + }; + + components.push(output); + } + + if self.num_files > 0 { + let output = if self.num_files > 1 { + format!("{} {}", self.num_files, "files") + } else { + format!("{} {}", self.num_files, "file") + }; + + components.push(output); + } + + if self.num_links > 0 { + let output = if self.num_links > 1 { + format!("{} {}", self.num_links, "links") + } else { + format!("{} {}", self.num_links, "link") + }; + + components.push(output); + } + + write!(f, "{}", components.join(", ")) + } +} diff --git a/src/render/tree/mod.rs b/src/render/tree/mod.rs index a63926c0..e4f81506 100644 --- a/src/render/tree/mod.rs +++ b/src/render/tree/mod.rs @@ -1,4 +1,5 @@ use crate::render::{context::Context, disk_usage::file_size::FileSize, order::Order, styles}; +use count::FileCount; use error::Error; use ignore::{WalkBuilder, WalkParallel}; use indextree::{Arena, NodeId}; @@ -16,6 +17,9 @@ use std::{ }; use visitor::{BranchVisitorBuilder, TraversalState}; +/// Operations to handle and display aggregate file counts based on their type. +mod count; + /// Errors related to traversal, [Tree] construction, and the like. pub mod error; @@ -245,6 +249,16 @@ impl Tree { descendant_id.detach(tree); } } + + fn compute_file_count(node_id: NodeId, tree: &Arena) -> FileCount { + let mut count = FileCount::default(); + + for child_id in node_id.children(tree) { + count.update(tree[child_id].get()); + } + + count + } } impl TryFrom<&Context> for WalkParallel { @@ -271,13 +285,25 @@ impl Display for Tree { let inner = self.inner(); let level = self.level(); let ctx = self.context(); + let show_count = ctx.count; + let mut file_count_data = vec![]; let mut descendants = root.descendants(inner).skip(1).peekable(); - let root_node = inner[root].get(); + let mut display_node = |node_id: NodeId, prefix: &str| -> fmt::Result { + let node = inner[node_id].get(); - root_node.display(f, "", ctx)?; - writeln!(f)?; + node.display(f, prefix, ctx)?; + + if show_count { + let count = Self::compute_file_count(node_id, inner); + file_count_data.push(count); + } + + writeln!(f) + }; + + display_node(root, "")?; let mut prefix_components = vec![""]; @@ -289,7 +315,7 @@ impl Display for Tree { let theme = if current_node.is_symlink() { styles::get_link_theme() } else { - styles::get_theme() + styles::get_tree_theme() }; let mut siblings = current_node_id.following_siblings(inner).skip(1).peekable(); @@ -305,8 +331,7 @@ impl Display for Tree { let prefix = current_prefix_components.join(""); if current_node.depth <= level { - current_node.display(f, &prefix, ctx)?; - writeln!(f)?; + display_node(current_node_id, &prefix)?; } if let Some(next_id) = descendants.peek() { @@ -326,6 +351,10 @@ impl Display for Tree { } } + if !file_count_data.is_empty() { + write!(f, "\n{}", FileCount::from(file_count_data))?; + } + Ok(()) } } diff --git a/src/render/tree/node/layout.rs b/src/render/tree/node/layout.rs new file mode 100644 index 00000000..4709bbdf --- /dev/null +++ b/src/render/tree/node/layout.rs @@ -0,0 +1,37 @@ +use crate::render::{context::Context, disk_usage::file_size::FileSize}; + +/// Simple struct to define location to put the `FileSize` while printing a `Node` +#[derive(Copy, Clone, Default)] +pub enum SizeLocation { + #[default] + Right, + Left, +} + +impl SizeLocation { + /// Returns a string to use when a node has no filesize, such as empty directories + pub fn default_string(self, ctx: &Context) -> String { + match self { + Self::Right => String::new(), + Self::Left => FileSize::empty_string(ctx), + } + } + + /// Given a [`FileSize`], style it in the expected way for its printing location + pub fn format(self, size: &FileSize) -> String { + match self { + Self::Right => format!("({})", size.format(false)), + Self::Left => size.format(true), + } + } +} + +impl From<&Context> for SizeLocation { + fn from(ctx: &Context) -> Self { + if ctx.size_left && !ctx.suppress_size { + Self::Left + } else { + Self::Right + } + } +} diff --git a/src/render/tree/node.rs b/src/render/tree/node/mod.rs similarity index 91% rename from src/render/tree/node.rs rename to src/render/tree/node/mod.rs index 31bbed89..365e76a8 100644 --- a/src/render/tree/node.rs +++ b/src/render/tree/node/mod.rs @@ -11,6 +11,7 @@ use ansi_term::Color; use ansi_term::Style; use ignore::DirEntry; use indextree::{Arena, Node as NodeWrapper, NodeId}; +use layout::SizeLocation; use lscolors::Style as LS_Style; use std::{ borrow::{Cow, ToOwned}, @@ -21,6 +22,9 @@ use std::{ path::{Path, PathBuf}, }; +/// For determining orientation of disk usage information for [Node]. +mod layout; + /// A node of [`Tree`] that can be created from a [DirEntry]. Any filesystem I/O and /// relevant system calls are expected to complete after initialization. A `Node` when `Display`ed /// uses ANSI colors determined by the file-type and [`LS_COLORS`]. @@ -224,10 +228,10 @@ impl Node { match size_loc { SizeLocation::Right => { - write!(f, "{prefix}{icon: { - write!(f, "{size} {prefix}{icon:)> for &Node { tree.get(node_id).map(NodeWrapper::get).unwrap() } } - -/// Simple struct to define location to put the `FileSize` while printing a `Node` -#[derive(Copy, Clone, Default)] -enum SizeLocation { - #[default] - Right, - Left, -} - -impl SizeLocation { - /// Returns a string to use when a node has no filesize, such as empty directories - fn default_string(self, ctx: &Context) -> String { - match self { - Self::Right => String::new(), - Self::Left => FileSize::empty_string(ctx), - } - } - - /// Given a [`FileSize`], style it in the expected way for its printing location - fn format(self, size: &FileSize) -> String { - match self { - Self::Right => format!("({})", size.format(false)), - Self::Left => size.format(true), - } - } -} - -impl From<&Context> for SizeLocation { - fn from(ctx: &Context) -> Self { - if ctx.size_left && !ctx.suppress_size { - Self::Left - } else { - Self::Right - } - } -} diff --git a/src/render/tree/report.rs b/src/render/tree/report.rs index 6272a25b..8a909db1 100644 --- a/src/render/tree/report.rs +++ b/src/render/tree/report.rs @@ -1,4 +1,4 @@ -use super::{node::Node, Tree}; +use super::{node::Node, FileCount, Tree}; use crate::render::disk_usage::{ file_size::{FileSize, HumanReadableComponents}, units::PrefixKind, @@ -10,11 +10,13 @@ use std::{ path::Path, }; +/// For a plain text output of disk usage information akin to `du`. pub struct Report<'a> { tree: &'a Tree, } impl<'a> Report<'a> { + /// Initializes a [Self] with a reference to [Tree]. pub const fn new(tree: &'a Tree) -> Self { Self { tree } } @@ -28,6 +30,8 @@ impl Display for Report<'_> { let max_depth = ctx.level().unwrap_or(usize::MAX); let dir = ctx.dir(); let prefix_kind = ctx.prefix; + let show_count = ctx.count; + let mut file_count_data = vec![]; let du_info = |node: &Node| { if ctx.human { @@ -50,6 +54,11 @@ impl Display for Report<'_> { let root_node = tree[root].get(); + if show_count { + let count = Tree::compute_file_count(root, tree); + file_count_data.push(count); + } + let total_du_width = root_node .file_size() .map_or_else(|| String::from("0"), |fs| format!("{}", fs.bytes)) @@ -84,6 +93,11 @@ impl Display for Report<'_> { for node_id in root.descendants(tree).skip(1) { let node = tree[node_id].get(); + if show_count { + let count = Tree::compute_file_count(node_id, tree); + file_count_data.push(count); + } + if node.depth > max_depth { continue; } @@ -106,6 +120,10 @@ impl Display for Report<'_> { writeln!(f, "{ft_iden} {du_info:>width_du_col$} {file}")?; } + if !file_count_data.is_empty() { + write!(f, "\n{}", FileCount::from(file_count_data))?; + } + Ok(()) } } diff --git a/tests/file_count.rs b/tests/file_count.rs new file mode 100644 index 00000000..45edda9c --- /dev/null +++ b/tests/file_count.rs @@ -0,0 +1,48 @@ +use indoc::indoc; + +mod utils; + +#[test] +fn file_count() { + assert_eq!( + utils::run_cmd(&["--count", "--sort", "name", "--no-config", "tests/data"]), + indoc!( + " + data (1.21 KiB) + ├─ dream_cycle (308 B) + │ └─ polaris.txt (308 B) + ├─ lipsum (446 B) + │ └─ lipsum.txt (446 B) + ├─ necronomicon.txt (83 B) + ├─ nemesis.txt (161 B) + ├─ nylarlathotep.txt (100 B) + └─ the_yellow_king (143 B) + └─ cassildas_song.md (143 B) + + 3 directories, 6 files" + ), + ) +} + +#[test] +fn file_count_report() { + assert_eq!( + utils::run_cmd(&["--count", "--report", "--sort", "name", "--no-config", "tests/data"]), + indoc!( + " + d 1241 B data + d 308 B dream_cycle + - 308 B dream_cycle/polaris.txt + d 446 B lipsum + - 446 B lipsum/lipsum.txt + - 83 B necronomicon.txt + - 161 B nemesis.txt + - 100 B nylarlathotep.txt + d 143 B the_yellow_king + - 143 B the_yellow_king/cassildas_song.md + + 3 directories, 6 files" + ), + ) +} +