count_files/
lib.rs

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

count_files/
lib.rs
