Feature/file count

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 <DISK_USAGE>    Print physical or logical file size [default: logical] [possible values: logical, physical]
   -g, --glob <GLOB>                Include or exclude files using glob patterns
       --iglob <IGLOB>              Include or exclude files using glob patterns; case insensitive

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<PathBuf>,
 

src/render/styles.rs

@@ -27,7 +27,7 @@ pub static LS_COLORS: OnceCell<LsColors> = 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<ThemesMap> = OnceCell::new();
+pub static TREE_THEME: OnceCell<ThemesMap> = 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)),

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<Vec<Self>> for FileCount {
+    fn from(data: Vec<Self>) -> 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(", "))
+    }
+}

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<Node>) -> 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(())
     }
 }

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

src/render/tree/node.rs → 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:<icon_padding$}{styled_name} {size}",)
+                write!(f, "{prefix}{icon:<icon_padding$}{styled_name} {size}")
             }
             SizeLocation::Left => {
-                write!(f, "{size} {prefix}{icon:<icon_padding$}{styled_name}",)
+                write!(f, "{size} {prefix}{icon:<icon_padding$}{styled_name}")
             }
         }
     }
@@ -356,39 +360,3 @@ impl From<(NodeId, &mut Arena<Self>)> 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
-        }
-    }
-}