Fix/bugfixes and code quality

src/classifier.rs

@@ -68,19 +68,29 @@ impl Classifier for PythonClassifier {
             return LineType::Comment;
         }
 
-        if trimmed.starts_with("\"\"\"") {
-            let count = line.matches("\"\"\"").count();
-            if count >= 2 {
+        if let Some(idx) = trimmed.find("\"\"\"") {
+            if trimmed.matches("\"\"\"").count() >= 2 {
+                // Single-line docstring. Check for code before or after.
+                let before = &trimmed[..idx];
+                let after = &trimmed[idx + 3..];
+                if !before.trim().is_empty() || !after.trim().is_empty() {
+                    return LineType::Pure;
+                }
                 return LineType::Docstring;
             } else {
                 self.in_triple_double = true;
                 return LineType::Docstring;
             }
         }
 
-        if trimmed.starts_with("'''") {
-            let count = line.matches("'''").count();
-            if count >= 2 {
+        if let Some(idx) = trimmed.find("'''") {
+            if trimmed.matches("'''").count() >= 2 {
+                // Single-line docstring. Check for code before or after.
+                let before = &trimmed[..idx];
+                let after = &trimmed[idx + 3..];
+                if !before.trim().is_empty() || !after.trim().is_empty() {
+                    return LineType::Pure;
+                }
                 return LineType::Docstring;
             } else {
                 self.in_triple_single = true;
@@ -133,9 +143,21 @@ impl Classifier for CStyleClassifier {
         if let Some(start_idx) = trimmed.find("/*") {
             if let Some(end_idx) = trimmed.find("*/") {
                 if end_idx > start_idx {
+                    // Single-line block comment. Check for code before or after.
+                    let before = &trimmed[..start_idx];
+                    let after = &trimmed[end_idx + 2..];
+                    if !before.trim().is_empty() || !after.trim().is_empty() {
+                        return LineType::Pure;
+                    }
                     return LineType::Comment;
                 }
             }
+            // Starts but doesn't end on same line. Check for code before /*.
+            let before = &trimmed[..start_idx];
+            if !before.trim().is_empty() {
+                self.in_block = true;
+                return LineType::Pure;
+            }
             self.in_block = true;
             return LineType::Comment;
         }
@@ -287,8 +309,11 @@ pub fn get_classifier(lang: Language) -> Box<dyn Classifier> {
         | Language::Swift
         | Language::Kotlin
         | Language::Scala
-        | Language::Css => Box::new(CStyleClassifier::new()),
-        Language::Shell | Language::PowerShell => Box::new(ShellClassifier),
+        | Language::Css
+        | Language::Rust => Box::new(CStyleClassifier::new()),
+        Language::Shell | Language::PowerShell | Language::Yaml | Language::Toml => {
+            Box::new(ShellClassifier)
+        }
         Language::Ruby => Box::new(RubyClassifier::new()),
         Language::Html | Language::Vue => Box::new(HtmlClassifier::new()),
         Language::Other => Box::new(DefaultClassifier),

src/files.rs

@@ -8,9 +8,10 @@ use std::path::Path;
 use walkdir::WalkDir;
 
 pub fn analyze_files(
+    paths: &[String],
     include: &[String],
     exclude: &[String],
-    reader: Option<Box<dyn BufRead>>, // For stdin support
+    reader: Option<Box<dyn BufRead>>,
 ) -> Result<Vec<FileStats>, std::io::Error> {
     let mut stats = Vec::new();
 
@@ -40,32 +41,31 @@ pub fn analyze_files(
         .filter_map(|p| Pattern::new(p).ok())
         .collect();
 
-    for entry in WalkDir::new(".").into_iter().flatten() {
-        let path = entry.path();
-        if path.is_dir() {
-            continue;
-        }
+    for root in paths {
+        for entry in WalkDir::new(root).into_iter().flatten() {
+            let path = entry.path();
+            if path.is_dir() {
+                continue;
+            }
 
-        // Convert path to relative string for matching
-        let path_str = path.to_string_lossy();
-        let clean_path = if let Some(stripped) = path_str.strip_prefix("./") {
-            stripped
-        } else {
-            &path_str
-        };
-
-        // Check excludes
-        if exclude_patterns.iter().any(|p| p.matches(clean_path)) {
-            continue;
-        }
+            let path_str = path.to_string_lossy();
+            let clean_path = if let Some(stripped) = path_str.strip_prefix("./") {
+                stripped
+            } else {
+                &path_str
+            };
 
-        // Check includes (at least one must match if we are strict, or default include is "**/*")
-        if !include_patterns.iter().any(|p| p.matches(clean_path)) {
-            continue;
-        }
+            if exclude_patterns.iter().any(|p| p.matches(clean_path)) {
+                continue;
+            }
+
+            if !include_patterns.iter().any(|p| p.matches(clean_path)) {
+                continue;
+            }
 
-        if let Ok(fs) = process_file(path) {
-            stats.push(fs);
+            if let Ok(fs) = process_file(path) {
+                stats.push(fs);
+            }
         }
     }
 

src/language.rs

@@ -21,6 +21,9 @@ pub enum Language {
     Shell,
     PowerShell,
     Vue,
+    Rust,
+    Yaml,
+    Toml,
     Other,
 }
 
@@ -45,12 +48,20 @@ impl Language {
             Some("sh") | Some("bash") | Some("zsh") => Language::Shell,
             Some("ps1") | Some("psm1") => Language::PowerShell,
             Some("vue") => Language::Vue,
+            Some("rs") => Language::Rust,
+            Some("yml") | Some("yaml") => Language::Yaml,
+            Some("toml") => Language::Toml,
             _ => {
                 // Check filename for special cases
-                match path.file_name().and_then(|n| n.to_str()) {
-                    Some("Dockerfile") => Language::Other, // Or maybe shell-like? keeping Other for now
-                    Some("Makefile") => Language::Other,
-                    _ => Language::Other,
+                let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
+                if name == "Dockerfile"
+                    || name.starts_with("Dockerfile.")
+                    || name == "Makefile"
+                    || name.ends_with(".mk")
+                {
+                    Language::Shell
+                } else {
+                    Language::Other
                 }
             }
         }
@@ -78,6 +89,9 @@ impl fmt::Display for Language {
             Language::Shell => "Shell",
             Language::PowerShell => "PowerShell",
             Language::Vue => "Vue",
+            Language::Rust => "Rust",
+            Language::Yaml => "YAML",
+            Language::Toml => "TOML",
             Language::Other => "Other",
         };
         write!(f, "{}", s)