1 files changed, 55 insertions, 34 deletions

diff --git a/src/main.rs b/src/main.rs
index 14ff4eb..737ddc1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -23,7 +23,6 @@ enum Error {
     UnknownCategory(OsString),
     NoCategories,
     DuplicateCategory(OsString),
-    BadPath(PathBuf),
     EmptyOption,
     InvalidLiteral(String),
 }
@@ -51,8 +50,6 @@ impl fmt::Display for Error {
                 write!(f, "no categories provided"),
             Error::DuplicateCategory(c) =>
                 write!(f, "tried to create a category twice: {:?}", c),
-            Error::BadPath(p) =>
-                write!(f, "failed to parse path: {:?}", p),
             Error::EmptyOption =>
                 write!(f, "multi-option variable had an empty option"),
             Error::InvalidLiteral(s) =>
@@ -287,45 +284,67 @@ impl SentenceSet {
     }
 }
 
-fn read_lines(path: &PathBuf) -> ProgResult<Vec<String>> {
-    let mut lines = vec![];
+struct ValidTextFile {
+    path: PathBuf,
+    stem: OsString,
+}
+
+impl ValidTextFile {
+    fn from_path(path: PathBuf) -> Option<Self> {
+        if !path.is_file() || path.extension()? != OsStr::new("txt") {
+            return None;
+        }
 
-    let text = fs::read_to_string(path)?;
-    for line in text.lines() {
-        lines.push(line.to_string());
+        let stem = path.file_stem()?.to_os_string();
+        Some(ValidTextFile {
+            path,
+            stem,
+        })
     }
 
-    Ok(lines)
-}
+    fn read_lines(&self) -> ProgResult<Vec<String>> {
+        let v = fs::read_to_string(&self.path)?
+            .lines()
+            .map(|s| s.to_string())
+            .collect();
+        Ok(v)
+    }
 
-// REVIEW: There has to be a better way to do this. Separate into multiple functions?
-fn read_files(dir_path: &str, unique: bool) -> ProgResult<(SentenceSet, CategorySet)> {
-    let mut opt_sentences: Option<SentenceSet> = None;
-    let mut categories = CategorySet::new(unique);
+    fn get_stem(&self) -> &OsStr {
+        &self.stem
+    }
+}
 
+fn find_text_files(dir_path: &str) -> ProgResult<Vec<ValidTextFile>> {
     let dir = fs::read_dir(dir_path)?;
+    let mut text_files = vec![];
+
     for entry in dir {
         let path_buf = entry?.path();
-        if path_buf.is_file() {
-            let extension = path_buf.extension()
-                .ok_or(Error::BadPath(path_buf.clone()))?;
-            if extension == OsStr::new("txt") {
-                let lines = read_lines(&path_buf)?;
-                let filename = path_buf.file_name()
-                    .ok_or(Error::BadPath(path_buf.clone()))?;
-                if filename == OsStr::new("sentences.txt") {
-                    if opt_sentences.is_none() {
-                        opt_sentences = Some(SentenceSet::new(lines, unique)?);
-                    } else {
-                        return Err(Error::MultipleSentences);
-                    }
-                } else {
-                    let cat_name = path_buf.file_stem()
-                        .expect("filename stem unwrapping osstr")
-                        .to_os_string();
-                    categories.add_category(cat_name, lines)?;
-                }
+        if let Some(text_file) = ValidTextFile::from_path(path_buf) {
+            text_files.push(text_file);
+        }
+    }
+
+    Ok(text_files)
+}
+
+fn load_files(files: &[ValidTextFile], unique: bool) -> ProgResult<(SentenceSet, CategorySet)> {
+    let mut opt_sentences: Option<SentenceSet> = None;
+    let mut categories = CategorySet::new(unique);
+
+    for file in files.iter() {
+        let stem = file.get_stem();
+        let lines = file.read_lines()?;
+
+        if stem == OsStr::new("sentences") {
+            if opt_sentences.is_none() {
+                opt_sentences = Some(SentenceSet::new(lines, unique)?);
+            } else {
+                return Err(Error::MultipleSentences);
             }
+        } else {
+            categories.add_category(stem.to_os_string(), lines)?;
         }
     }
 
@@ -342,7 +361,9 @@ fn crash(e: Error) -> ! {
 
 fn run() -> ProgResult<()> {
     let args = Args::parse()?;
-    let (mut sentences, mut categories) = read_files(&args.directory, args.unique)?;
+
+    let files = find_text_files(&args.directory)?;
+    let (mut sentences, mut categories) = load_files(files.as_slice(), args.unique)?;
 
     for _i in 0..args.n_sentences {
         let sentence = sentences.random_sentence()?;