1 files changed, 84 insertions, 26 deletions

diff --git a/src/main.rs b/src/main.rs
index ee94d6e..913bcbb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,5 @@
-use std::io::stdin;
-use std::io::BufRead;
-use std::io::BufReader;
+use std::env::args;
+use std::io::{stdin, BufRead, BufReader};
 
 /// Words that will not be capitalized unless they are at the start or end of a title.
 const EXCEPTIONS: &[&str] = &[
@@ -18,46 +17,80 @@ const EXCEPTIONS: &[&str] = &[
 /// - Words are not capitalized if they are EXCEPTIONS, unless...
 /// - Words are either the first or the last word.
 fn main() {
+    let mut args = args();
+    args.next(); // ignore self arg
+    let mut ignorables = Vec::new();
+
+    while let Some(arg) = args.next() {
+        dbg!(&arg);
+        match arg.as_str() {
+            "--ignore" | "-i" => {
+                ignorables = args.collect();
+                break;
+            }
+            "--help" | "-h" => {
+                println!("Title-Case Tool.\n-i | --ignore [word...]\tIgnores a list of words that should not be altered.");
+            }
+            unrecognized => {
+                panic!(format!("Unrecognized option: {}", unrecognized));
+            }
+        }
+    }
+
     for line in BufReader::new(stdin()).lines() {
-        println!("{}", capitalize_line(&line.expect("IO error")));
+        let line = line.expect("IO error");
+        println!("{}", correct_line(&line, &ignorables));
     }
 }
 
 /// Capitalizes a single line (title) based on placement and exceptions.
-fn capitalize_line(line: &str) -> String {
-    let mut line = line.split_whitespace().map(str::to_lowercase).peekable();
-    let mut result: Vec<String> = Vec::new();
+fn correct_line(line: &str, ignorables: &[String]) -> String {
+    let mut line = line.split_whitespace().peekable();
+    let mut result = Vec::new();
 
-    // handle first word
-    let first = line.next().expect("No words on line");
-    result.push(capitalize_word(&first));
+    if let Some(first) = line.next() {
+        // always correct the first word
+        result.push(correct_word(first, true, ignorables));
+    } else {
+        // handle empty lines
+        return String::new();
+    }
 
+    // handle every other word
     while let Some(word) = line.next() {
-        let word = if line.peek().is_none() || !EXCEPTIONS.contains(&word.as_str()) {
-            // capitalize words that are last or aren't in exceptions
-            capitalize_word(&word)
-        } else {
-            // ignore the rest
-            word
-        };
-
-        result.push(word);
+        result.push(correct_word(word, line.peek().is_none(), ignorables));
     }
 
     // join words into title
     result.join(" ")
 }
 
-/// Capitalizes a single word.
+/// Corrects a single word based on exceptions and ignorables.
+fn correct_word(word: &str, first_or_last: bool, ignorables: &[String]) -> String {
+    if ignorables.iter().any(|s| s == word) {
+        // leave ignorables untouched
+        word.to_owned()
+    } else if first_or_last || !EXCEPTIONS.iter().any(|s| **s == word.to_lowercase()) {
+        // capitalize words that are first, last, or aren't exceptions
+        capitalize_word(&word.to_lowercase())
+    } else {
+        // lowercase all other words
+        word.to_lowercase()
+    }
+}
+
+// correct with just ignorables
+// correct with ignorables and exceptions
+
+/// Capitalizes the first letter of a single word.
 fn capitalize_word(word: &str) -> String {
     let mut chars = word.chars();
-    let first = chars.next().unwrap();
-    format!("{}{}", first.to_uppercase(), chars.collect::<String>())
+    let first = chars.next().unwrap_or_default();
+    first.to_uppercase().chain(chars).collect()
 }
 
 #[cfg(test)]
 mod tests {
-    // Note this useful idiom: importing names from outer (for mod tests) scope.
     use super::*;
 
     #[test]
@@ -67,14 +100,39 @@ mod tests {
     }
 
     #[test]
-    fn test_capitalize_line() {
+    fn test_correct_word() {
+        assert_eq!(correct_word("Like", false, &[]), String::from("like"));
+        assert_eq!(correct_word("like", false, &[]), String::from("like"));
+        assert_eq!(correct_word("like", true, &[]), String::from("Like"));
+        assert_eq!(correct_word("liKe", false, &[]), String::from("like"));
         assert_eq!(
-            capitalize_line("this is a test"),
+            correct_word("computers", false, &[]),
+            String::from("Computers")
+        );
+    }
+
+    #[test]
+    fn test_correct_line() {
+        assert_eq!(
+            correct_line("this is a test", &[]),
             String::from("This Is a Test")
         );
         assert_eq!(
-            capitalize_line("similEs: lIke Or As"),
+            correct_line("similEs: lIke Or As", &[]),
             String::from("Similes: like or As")
         );
     }
+
+    #[test]
+    fn test_proper_nouns_acronyms() {
+        assert_eq!(
+            correct_line("FreeBSD", &[String::from("FreeBSD")]),
+            String::from("FreeBSD")
+        );
+    }
+
+    #[test]
+    fn test_empty_lines() {
+        assert_eq!(correct_line("", &[]), String::from(""));
+    }
 }