path: root/src/main.rs

use std::collections::HashSet;
use std::env::args;
use std::io::{prelude::*, stdin, stdout, BufReader, BufWriter};

const HELP_MSG: &str = "Title-Case Tool\n\
        -h | --help\t\t\tPrints this message\n\
        -v | --version\t\t\tPrints version number\n\
        -i | --ignore [word...]\t\tA list of words that should not be corrected";

const VERSION_MSG: &str = env!("CARGO_PKG_VERSION");

/// Words that will not be capitalized unless they are at the start or end of a title.
const EXCEPTIONS: &[&str] = &[
    "vs", "above", "across", "against", "at", "between", "by", "along", "among", "down", "in",
    "once", "around", "of", "off", "on", "to", "with", "before", "behind", "below", "beneath",
    "down", "from", "near", "toward", "upon", "within", "a", "an", "the", "and", "but", "for",
    "or", "nor", "so", "till", "when", "yet", "that", "than", "in", "into", "like", "onto", "over",
    "past", "as", "if",
];

/// Here are the rules:
///
/// - Every line is its own title.
/// - All words are capitalized, except...
/// - Words are not capitalized if they are EXCEPTIONS, unless...
/// - Words are either the first or the last word.
fn main() {
    // gather args
    let mut args = args();
    args.next(); // ignore self arg
    let mut ignorables = HashSet::new();

    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--ignore" | "-i" => {
                ignorables = args.collect();
                break;
            }
            "--help" | "-h" => {
                println!("{}", HELP_MSG);
                return;
            }
            "--version" | "-v" => {
                println!("{}", VERSION_MSG);
                return;
            }
            unrecognized => {
                panic!(format!("Unrecognized option: {}", unrecognized));
            }
        }
    }

    // gather exceptions
    let exceptions: HashSet<&'static str> = EXCEPTIONS.iter().cloned().collect();

    // read input, correct lines, write output
    let mut writer = BufWriter::new(stdout());
    for line in BufReader::new(stdin()).lines() {
        let line = line.expect("IO error");
        writeln!(
            &mut writer,
            "{}",
            correct_line(&line, &ignorables, &exceptions)
        )
        .expect("Failed to write output");
    }
    writer.flush().expect("Failed to write output");
}

/// Corrects a single line (title) based on placement and exceptions.
fn correct_line(
    line: &str,
    ignorables: &HashSet<String>,
    exceptions: &HashSet<&'static str>,
) -> String {
    let mut line = line.split_whitespace().peekable();
    let mut result = String::new();

    if let Some(first) = line.next() {
        // always correct the first word
        result.push_str(&correct_word(first, true, ignorables, exceptions));
    } else {
        // skip empty lines
        return result;
    }

    // handle every other word
    while let Some(word) = line.next() {
        result.push(' ');
        result.push_str(&correct_word(
            word,
            line.peek().is_none(),
            ignorables,
            exceptions,
        ));
    }

    result
}

/// Corrects a single word based on exceptions and ignorables.
fn correct_word(
    word: &str,
    first_or_last: bool,
    ignorables: &HashSet<String>,
    exceptions: &HashSet<&'static str>,
) -> String {
    if ignorables.iter().any(|s| s == word) {
        // leave ignorables untouched
        word.to_owned()
    } else if first_or_last || !exceptions.contains(&word.to_lowercase().as_str()) {
        // capitalize words that are first, last, or aren't exceptions
        capitalize_word(&word.to_lowercase())
    } else {
        // lowercase all other words
        word.to_lowercase()
    }
}

/// Capitalizes the first letter of a single word.
fn capitalize_word(word: &str) -> String {
    let mut chars = word.chars();

    if let Some(first) = chars.next() {
        // uppercase first character and append the rest
        first.to_uppercase().chain(chars).collect()
    } else {
        // handle empty words
        String::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_capitalize_word() {
        assert_eq!(capitalize_word("word"), String::from("Word"));
        assert_eq!(capitalize_word("as"), String::from("As"));
    }

    #[test]
    fn test_correct_word() {
        assert_eq!(
            correct_word(
                "Like",
                false,
                &HashSet::new(),
                &EXCEPTIONS.iter().cloned().collect()
            ),
            String::from("like")
        );
        assert_eq!(
            correct_word(
                "like",
                false,
                &HashSet::new(),
                &EXCEPTIONS.iter().cloned().collect()
            ),
            String::from("like")
        );
        assert_eq!(
            correct_word(
                "like",
                true,
                &HashSet::new(),
                &EXCEPTIONS.iter().cloned().collect()
            ),
            String::from("Like")
        );
        assert_eq!(
            correct_word(
                "liKe",
                false,
                &HashSet::new(),
                &EXCEPTIONS.iter().cloned().collect()
            ),
            String::from("like")
        );
        assert_eq!(
            correct_word("computers", false, &HashSet::new(), &HashSet::new()),
            String::from("Computers")
        );
    }

    #[test]
    fn test_correct_line() {
        assert_eq!(
            correct_line(
                "this is a test",
                &HashSet::new(),
                &EXCEPTIONS.iter().cloned().collect()
            ),
            String::from("This Is a Test")
        );
        assert_eq!(
            correct_line(
                "similEs: lIke Or As",
                &HashSet::new(),
                &EXCEPTIONS.iter().cloned().collect()
            ),
            String::from("Similes: like or As")
        );
    }

    #[test]
    fn test_proper_nouns_acronyms() {
        let mut ignorables = HashSet::new();
        ignorables.insert(String::from("FreeBSD"));
        assert_eq!(
            correct_line("FreeBSD", &ignorables, &HashSet::new()),
            String::from("FreeBSD")
        );
    }

    #[test]
    fn test_empty_lines() {
        assert_eq!(
            correct_line("", &HashSet::new(), &HashSet::new()),
            String::from(""),
        );
    }
}