Implementing the Tree command in Rust, part 2: Printing Trees
In Part 1, we saw how to walk directory trees, recursively using fs::read_dir to construct an in-memory tree of FileNodes. In Part 2, we’ll implement the rest of the core of the tree command: printing the directory tree with Box Drawing characters.
Let’s take a look at some output from tree:
. ├── alloc.rs ├── ascii.rs ├── os │ ├── wasi │ │ ├── ffi.rs │ │ ├── mod.rs ➊ │ │ └── net ➋ │ │ └── mod.rs │ └── windows │ ├── ffi.rs ➌ │ ├── fs.rs │ ├── io │ │ └── tests.rs │ ├── mod.rs │ └── thread.rs ├── personality │ ├── dwarf │ │ ├── eh.rs │ │ ├── mod.rs │ │ └── tests.rs │ ├── emcc.rs │ └── gcc.rs └── personality.rs
The first thing that we notice is that most entries at any level, such as ➊, are preceded by ├──, while the last entry, ➋, is preceded by └──. This article about building a directory tree generator in Python calls them the tee and elbow connectors, and I’m going to use that terminology.
The second thing we notice is that there are multiple prefixes before the connectors, either │ (pipe) or (space), one prefix for each level. The rule is that children of a last entry, such as os/windows ➌, get the space prefix, while children of other entries, such as os/wasi or personality, get the pipe prefix.
For both connectors and prefixes, the last entry at a particular level gets special treatment.
The print_tree function
A classic technique with recursion is to create a pair of functions: an outer public function that calls a private helper function with the initial set of parameters to visit recursively.
Our print_tree function uses an inner visit function to recursively do almost all of the work.
pub fn print_tree(root: &str, dir: &Directory) { const OTHER_CHILD: &str = "│ "; // prefix: pipe const OTHER_ENTRY: &str = "├── "; // connector: tee const FINAL_CHILD: &str = " "; // prefix: no more siblings const FINAL_ENTRY: &str = "└── "; // connector: elbow println!("{}", root); ➊ let (d, f) = visit(dir, ""); println!("\n{} directories, {} files", d, f); fn visit(node: &Directory, prefix: &str) -> (usize, usize) { ➋ let mut dirs: usize = 1; // counting this directory ➌ let mut files: usize = 0; let mut count = node.entries.len(); ➍ for entry in &node.entries { count -= 1; let connector = if count == 0 { FINAL_ENTRY } else { OTHER_ENTRY }; ➎ match entry { FileTree::DirNode(sub_dir) => { ➏ println!("{}{}{}", prefix, connector, sub_dir.name); let new_prefix = format!( ➐ "{}{}", prefix, if count == 0 { FINAL_CHILD } else { OTHER_CHILD } ); let (d, f) = visit(&sub_dir, &new_prefix); ➑ dirs += d; files += f; } FileTree::LinkNode(symlink) => { println!( "{}{}{} -> {}", prefix, connector, symlink.name, symlink.target); files += 1; } FileTree::FileNode(file) => { println!("{}{}{}", prefix, connector, file.name); files += 1; } } } (dirs, files) ➒ } }
- The outer function, print_tree, simply prints the name of the root node on a line by itself; calls the inner visit function with the dir node and an empty prefix; and finally prints the number of directories and files visited. This is for compatibility with the output of tree.
- The inner visit function takes two parameters: node, a Directory, and prefix, a string which is initially empty.
- Keep track of the number of dirs and files seen at this level and in sub-directories.
- We count downwards from the number of entries in this directory to zero. When count is zero, we are on the last entry, which gets special treatment.
- Compute the connector, └── (elbow) for the last entry; ├── (tee) otherwise.
- Match the FileTree::DirNode variant and destructure the value into sub_dir, a &Directory.
- Before recursively visiting a sub-directory, we compute a new prefix, by appending the appropriate sub-prefix to the current prefix. If there are further entries (count > 0), the sub-prefix for the current level is │ (pipe); otherwise, it’s (spaces).
- Call visit recursively, then add to the running totals of dirs and files.
- visit returns a tuple of the counts of directories and files that were recursively visited.
One subtlety that is not obvious from the above is that OTHER_CHILD actually contains two non-breaking spaces:
const OTHER_CHILD: &str = "│\u{00A0}\u{00A0} "; // prefix: pipe
This is for compatibility with the output of tree:
$ diff <(cargo run -q -- ./tests) <(tree ./tests) && echo "no difference" no difference
Using process substitution to generate two different inputs for diff.
The main function
Let’s tie it all together.
fn main() -> io::Result<()> { let root = env::args().nth(1).unwrap_or(".".to_string()); ➊ let dir: Directory = dir_walk( ➋ &PathBuf::from(root.clone()), ➌ is_not_hidden, sort_by_name)?; ➍ print_tree(&root, &dir); ➎ Ok(()) ➏ }
- The simplest possible way to get a single, optional command-line argument. If omitted, we default to ., the current directory. For more sophisticated argument parsing, we could use Clap.
- Use dir_walk from Part 1 to recursively build a directory of FileTree nodes.
- Create a PathBuf from root, a string; clone is needed because PathBuf::from takes ownership of the string buffer. Use the is_not_hidden filter and the sort_by_name comparator from Part 1.
- The postfix question mark operator, ?, is used to propagate errors.
- Let print_tree draw the diagram.
- Return the Ok unit result to indicate success.
Resources
- Official tree source: The actual source for tree, written in old-school C.
- Draw a Tree Structure With Only CSS: Use CSS to draw links in nested, unordered lists.
- Build a Python Directory Tree Generator for the Command Line.
- Kevin Newton has implemented Tree in Multiple Languages.
- Tre is a modern alternative to tree in Rust.