diff --git a/src/render/context/mod.rs b/src/render/context/mod.rs index 50cb365a..6874e6c3 100644 --- a/src/render/context/mod.rs +++ b/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, + + #[arg(short = 'D', long)] + /// Only show directories + pub dirs_only: bool, } impl Context { diff --git a/src/render/tree/mod.rs b/src/render/tree/mod.rs index 23998db1..558e85e0 100644 --- a/src/render/tree/mod.rs +++ b/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, 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> = 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, 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) { + 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 { diff --git a/tests/dirs_only.rs b/tests/dirs_only.rs new file mode 100644 index 00000000..96f89cdf --- /dev/null +++ b/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." + ) +}