Add a file size alias as CLI input (#6)

* Implement numeric `min_file_size` * Implement file size via aliases
yury-fedotov · Jul 31, 2024 · b16bea7 · b16bea7
1 parent 2cb7a31
commit b16bea7
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 4 deletions.
diff --git a/src/cli.rs b/src/cli.rs
@@ -1,7 +1,27 @@
 use clap::{Arg, Command};
 
+const SIZE_ALIASES: &[(&str, usize)] = &[
+    ("blank", 1024),               // 1 KB
+    ("config", 10 * 1024),         // 10 KB
+    ("code", 100 * 1024),          // 100 KB
+    ("excel", 1024 * 1024),        // 1 MB
+    ("document", 5 * 1024 * 1024), // 5 MB
+    ("image", 10 * 1024 * 1024),   // 10 MB
+    ("gif", 20 * 1024 * 1024),     // 20 MB
+    ("audio", 50 * 1024 * 1024),   // 50 MB
+    ("video", 500 * 1024 * 1024),  // 500 MB
+    ("large", 1024 * 1024 * 1024), // 1 GB
+];
+
+fn get_size_by_alias(alias: &str) -> Option<usize> {
+    SIZE_ALIASES
+        .iter()
+        .find_map(|&(key, size)| if key == alias { Some(size) } else { None })
+}
+
 pub struct CLIArgs {
     pub directory: String,
+    pub min_file_size: usize,
     pub n_top: usize,
 }
 
@@ -15,6 +35,12 @@ pub fn parse_args() -> CLIArgs {
                 .required(true)
                 .index(1),
         )
+        .arg(
+            Arg::new("min_file_size")
+                .help("Minimum file size to consider (alias)")
+                .long("min_file_size")
+                .default_value("document"),
+        )
         .arg(
             Arg::new("n_top")
                 .help("Number of top files to return based on size")
@@ -28,10 +54,19 @@ pub fn parse_args() -> CLIArgs {
         .get_one::<String>("directory")
         .expect("Directory argument missing")
         .clone();
+    let size_alias = matches
+        .get_one::<String>("min_file_size")
+        .expect("Size argument missing");
     let n_top = matches
         .get_one::<String>("n_top")
         .map(|s| s.parse().unwrap_or(5)) // Parse the value and default to 5 on error
         .unwrap_or(5);
 
-    CLIArgs { directory, n_top }
+    let min_file_size = get_size_by_alias(size_alias.as_str()).unwrap_or(1024 * 1024 * 1024); // Default to 1 GB if alias not found
+
+    CLIArgs {
+        directory,
+        min_file_size,
+        n_top,
+    }
 }
diff --git a/src/main.rs b/src/main.rs
@@ -14,7 +14,7 @@ fn main() {
 
     let args = parse_args();
 
-    let traverval_output = traverse_directory(&args.directory);
+    let traverval_output = traverse_directory(&args.directory, &args.min_file_size);
     let file_count = traverval_output.file_infos.len();
     let largest_files = get_largest_files(&traverval_output.file_infos, &args.n_top);
     let dir_count = traverval_output.dir_count;

diff --git a/src/traversal.rs b/src/traversal.rs
@@ -8,7 +8,7 @@ pub struct DirectoryTraversalOutput {
     pub max_depth: usize,
 }
 
-pub fn traverse_directory(dir: &str) -> DirectoryTraversalOutput {
+pub fn traverse_directory(dir: &str, min_file_size: &usize) -> DirectoryTraversalOutput {
     let mut dir_count = 1;
     let mut max_depth = 0;
     let file_infos: Vec<FileInfo> = WalkBuilder::new(dir)
@@ -30,7 +30,7 @@ pub fn traverse_directory(dir: &str) -> DirectoryTraversalOutput {
             if entry.file_type().map_or(false, |ft| ft.is_file()) {
                 let path = entry.into_path();
                 let size = get_file_size(&path);
-                let hash = if size > 3 * 1024 * 1024 {
+                let hash = if size > *min_file_size as u64 {
                     // 3MB in bytes
                     match calculate_hash(&path) {
                         Ok(hash) => hash,