count_files/
lib.rs

1use clap::Parser;
2use comfy_table::Table;
3use indicatif::{HumanBytes, HumanDuration, ProgressBar, ProgressStyle};
4use std::collections::HashMap;
5use std::error::Error;
6use std::fs;
7use std::path::Path;
8use std::time::Instant;
9
10/// Counting files in a directory.
11#[derive(Parser)]
12#[clap(author, version, about, long_about = None)]
13pub struct Args {
14    /// Target directory to be scanned.
15    pub target_path: String,
16}
17
18// A struct to record the file count number and the total storage size.
19struct Counter {
20    count: u64,
21    storage_size: u64,
22}
23
24impl Counter {
25    // Create a new Counter struct.
26    pub fn new(count: u64, storage_size: u64) -> Counter {
27        Counter {
28            count,
29            storage_size,
30        }
31    }
32
33    // Update the counter.
34    pub fn update(&mut self, count: i64, storage_size: i64) {
35        self.count = (self.count as i64 + count) as u64;
36        self.storage_size = (self.storage_size as i64 + storage_size) as u64;
37    }
38}
39
40// Scan the target path and count all the files.
41fn scan(
42    path: &Path,
43    record: &mut HashMap<String, Counter>,
44    pb: ProgressBar,
45) -> Result<(), Box<dyn Error>> {
46    // Tell the user where are we now.
47    pb.set_message(path.to_str().unwrap().to_string());
48
49    // Loop the entries.
50    let entries = fs::read_dir(path)?;
51    for entry in entries {
52        let entry = entry?;
53        let path = entry.path();
54
55        // The entry is a directory or a file?
56        if path.is_dir() {
57            if let Err(e) = scan(path.as_path(), record, pb.clone()) {
58                println!("WARNING: {}. Skip {}", e, path.to_str().unwrap());
59            }
60        } else if let Some(extension) = path.extension() {
61            let extension = extension.to_str().unwrap().to_string();
62            let counter = record
63                .entry(extension)
64                .or_insert_with(|| Counter::new(0, 0));
65            // Get the size of the file in bytes.
66            let mut file_size: i64 = 0;
67            if let Ok(attribute) = fs::metadata(&path) {
68                file_size = attribute.len() as i64;
69            }
70
71            // Update the counter.
72            counter.update(1, file_size);
73        }
74    }
75    Ok(())
76}
77
78// Print the counting result in a table.
79fn print_to_screen(record: &HashMap<String, Counter>) {
80    // Sort the result by file count.
81    let mut record: Vec<(&String, &Counter)> = record.iter().collect();
82    record.sort_by(|a, b| b.1.count.cmp(&a.1.count));
83
84    // Create the result table.
85    let mut table = Table::new();
86    table.set_header(vec!["File type", "Count", "Total size"]);
87    for (ext, counter) in record {
88        table.add_row(vec![
89            ext,
90            &counter.count.to_string(),
91            &HumanBytes(counter.storage_size).to_string(),
92        ]);
93    }
94
95    // Align the numbers to right.
96    for i in 1..3 {
97        if let Some(column) = table.get_column_mut(i) {
98            column.set_cell_alignment(comfy_table::CellAlignment::Right)
99        }
100    }
101    println!("{table}");
102}
103
104// Count all the files.
105pub fn run(config: &Args) -> Result<(), Box<dyn Error>> {
106    // Use a hashmap to record different files count.
107    let mut record: HashMap<String, Counter> = HashMap::new();
108    let target_path = Path::new(&config.target_path);
109
110    // Setup a progress bar.
111    let spinner_style = ProgressStyle::default_spinner()
112        .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ")
113        .template("{prefix:.bold.dim} {spinner} {wide_msg}");
114    let pb = ProgressBar::new_spinner();
115    pb.set_style(spinner_style);
116    pb.set_prefix("Scanning ");
117
118    // Setup a duration meter.
119    let started = Instant::now();
120
121    // Let the party begin.
122    scan(target_path, &mut record, pb.clone())?;
123
124    // Post process
125    pb.finish_and_clear();
126    println!("Done in {}.", HumanDuration(started.elapsed()));
127    print_to_screen(&record);
128    Ok(())
129}
count_files/lib.rs

count_files/
lib.rs
