Implement dirs-only option

src/render/context/mod.rs

@@ -113,6 +113,10 @@ pub struct Context {
     #[arg(long)]
     /// Print completions for a given shell to stdout
     pub completions: Option<clap_complete::Shell>,
+
+    #[arg(short = 'D', long)]
+    /// Only show directories
+    pub dirs_only: bool,
 }
 
 impl Context {

src/render/tree/mod.rs

@@ -70,7 +70,7 @@ impl Tree {
 
     /// Parallel traversal of the root directory and its contents. Parallel traversal relies on
     /// `WalkParallel`. Any filesystem I/O or related system calls are expected to occur during
-    /// parallel traversal; post-processing post-processing of all directory entries should
+    /// parallel traversal; post-processing of all directory entries should
     /// be completely CPU-bound.
     fn traverse(ctx: &Context) -> TreeResult<(Arena<Node>, NodeId)> {
         let walker = WalkParallel::try_from(ctx)?;
@@ -83,7 +83,7 @@ impl Tree {
                 // Key represents path of parent directory and values represent children.
                 let mut branches: HashMap<PathBuf, Vec<NodeId>> = HashMap::new();
 
-                // Set used to prevent double counting hard-links in the same file-tree hiearchy.
+                // Set used to prevent double counting hard-links in the same file-tree hierarchy.
                 let mut inodes = HashSet::new();
 
                 let mut root_id = None;
@@ -131,6 +131,10 @@ impl Tree {
                     Self::prune_directories(root, &mut tree);
                 }
 
+                if ctx.dirs_only {
+                    Self::prune_files(root, &mut tree);
+                }
+
                 Ok::<(Arena<Node>, NodeId), Error>((tree, root))
             });
 
@@ -212,6 +216,25 @@ impl Tree {
             node_id.remove_subtree(tree)
         }
     }
+
+    /// Function to remove files.
+    fn prune_files(root_id: NodeId, tree: &mut Arena<Node>) {
+        let mut to_prune = vec![];
+
+        for node_id in root_id.descendants(tree) {
+            let node = Node::get(node_id, tree).unwrap();
+
+            if !node.is_dir() {
+                if node_id.children(tree).peekable().peek().is_none() {
+                    to_prune.push(node_id);
+                }
+            }
+        }
+
+        for node_id in to_prune {
+            node_id.remove_subtree(tree)
+        }
+    }
 }
 
 impl TryFrom<&Context> for WalkParallel {

tests/dirs_only.rs

@@ -0,0 +1,18 @@
+use indoc::indoc;
+
+mod utils;
+
+#[test]
+fn level() {
+    assert_eq!(
+        utils::run_cmd(&["--sort", "name", "--no-config", "--dirs-only", "tests/data"]),
+        indoc!(
+            "
+            data (1.21 KiB)
+            ├─ dream_cycle (308 B)
+            ├─ lipsum (446 B)
+            └─ the_yellow_king (143 B)"
+        ),
+        "Failed to print dirs-only."
+    )
+}