1 files changed, 75 insertions, 64 deletions

diff --git a/src/files.rs b/src/files.rs
index 2512fa5..81bdd25 100644
--- a/src/files.rs
+++ b/src/files.rs
@@ -1,9 +1,11 @@
 use std::{
-    fs::{self, DirEntry, File, OpenOptions},
+    fs::{self, DirEntry, File, OpenOptions, ReadDir},
     io,
     path::{Path, PathBuf},
 };
 
+use anyhow::{Context, Error, Result};
+
 use crate::actions::ActionOptions;
 
 pub struct Locations {
@@ -13,61 +15,68 @@ pub struct Locations {
 }
 
 impl Locations {
-    pub fn get_repository_paths(&self) -> Result<Vec<RepositoryPaths>, ()> {
-        let repository_entries = fs::read_dir(&self.repository_path).map_err(|_| ())?;
-        let history_entries = fs::read_dir(&self.ka_files_path).map_err(|_| ())?;
-
-        let tracked_paths = repository_entries
-            .map(Result::unwrap)
-            .filter(|entry| entry.path() != self.ka_path)
-            .flat_map(Self::flatten_directories)
-            .map(|entry| {
-                let path = entry.path();
-                RepositoryPaths::from_tracked(&self, &path)
-            });
-
-        let deleted_paths = history_entries
-            .map(Result::unwrap)
-            .flat_map(Self::flatten_directories)
-            .filter_map(|entry| {
-                let path = entry.path();
-                let file = RepositoryPaths::from_history(&self, &path);
-                match file {
-                    RepositoryPaths::Deleted { .. } => Some(file),
-                    RepositoryPaths::Tracked { .. } => None,
-                    _ => unreachable!(),
-                }
-            });
+    pub fn get_repository_paths(&self) -> Result<Vec<RepositoryPaths>, Error> {
+        let repository_entries =
+            fs::read_dir(&self.repository_path).context("Failed reading tracked file entries.")?;
+        let history_entries =
+            fs::read_dir(&self.ka_files_path).context("Failed reading history file entries.")?;
+
+        let tracked_paths = Self::walk_directory(repository_entries, &|entry| {
+            let file_path = entry.path();
+            if file_path != self.ka_path {
+                RepositoryPaths::from_tracked(&self, &file_path).ok()
+            } else {
+                None
+            }
+        })?;
+
+        let deleted_paths = Self::walk_directory(history_entries, &|entry| {
+            let file_path = entry.path();
+            let file = RepositoryPaths::from_history(&self, &file_path).ok()?;
+            match file {
+                RepositoryPaths::Deleted { .. } => Some(file),
+                RepositoryPaths::Tracked { .. } => None,
+                _ => unreachable!(),
+            }
+        })?;
 
-        let all_paths = tracked_paths.chain(deleted_paths);
+        let mut all_paths = tracked_paths;
+        all_paths.extend(deleted_paths);
 
-        Ok(all_paths.collect())
+        Ok(all_paths)
     }
 
-    pub fn tracked_from_history(&self, history_file_path: &Path) -> PathBuf {
-        self.repository_path
-            .join(history_file_path.strip_prefix(&self.ka_files_path).unwrap())
+    pub fn tracked_from_history(&self, history_file_path: &Path) -> Result<PathBuf> {
+        let raw_path = history_file_path.strip_prefix(&self.ka_files_path)?;
+        Ok(self.repository_path.join(raw_path))
     }
 
-    pub fn history_from_tracked(&self, tracked_file_path: &Path) -> PathBuf {
-        self.ka_files_path.join(
-            tracked_file_path
-                .strip_prefix(&self.repository_path)
-                .unwrap(),
-        )
+    pub fn history_from_tracked(&self, tracked_file_path: &Path) -> Result<PathBuf> {
+        let raw_path = tracked_file_path.strip_prefix(&self.repository_path)?;
+        Ok(self.ka_files_path.join(raw_path))
     }
 
-    fn flatten_directories(entry: DirEntry) -> Vec<DirEntry> {
-        let file_type = entry.file_type().expect("Could not read a file type.");
-        if file_type.is_dir() {
-            fs::read_dir(entry.path())
-                .expect("Could not read nested directory.")
-                .map(Result::unwrap)
-                .flat_map(Self::flatten_directories)
-                .collect()
-        } else {
-            vec![entry]
+    fn walk_directory(
+        directory: ReadDir,
+        filter_map: &dyn Fn(DirEntry) -> Option<RepositoryPaths>,
+    ) -> Result<Vec<RepositoryPaths>> {
+        let mut entries = Vec::new();
+
+        for res in directory {
+            let entry = res?;
+            let file_type = entry.file_type()?;
+            if file_type.is_dir() {
+                let nested_directory = fs::read_dir(entry.path())?;
+                let nested_paths = Self::walk_directory(nested_directory, filter_map)?;
+                entries.extend(nested_paths);
+            } else {
+                if let Some(paths) = filter_map(entry) {
+                    entries.push(paths);
+                }
+            }
         }
+
+        Ok(entries)
     }
 }
 
@@ -91,9 +100,9 @@ pub enum RepositoryPaths {
 }
 
 impl RepositoryPaths {
-    pub fn from_history(locations: &Locations, history_file_path: &Path) -> Self {
-        let tracked_path = locations.tracked_from_history(history_file_path);
-        if !tracked_path.exists() {
+    pub fn from_history(locations: &Locations, history_file_path: &Path) -> Result<Self> {
+        let tracked_path = locations.tracked_from_history(history_file_path)?;
+        Ok(if !tracked_path.exists() {
             RepositoryPaths::Deleted(FileDeleted {
                 history_path: history_file_path.to_path_buf(),
             })
@@ -102,12 +111,12 @@ impl RepositoryPaths {
                 history_path: history_file_path.to_path_buf(),
                 tracked_path,
             })
-        }
+        })
     }
 
-    pub fn from_tracked(locations: &Locations, tracked_file_path: &Path) -> Self {
-        let history_path = locations.history_from_tracked(tracked_file_path);
-        if !history_path.exists() {
+    pub fn from_tracked(locations: &Locations, tracked_file_path: &Path) -> Result<Self> {
+        let history_path = locations.history_from_tracked(tracked_file_path)?;
+        Ok(if !history_path.exists() {
             RepositoryPaths::Untracked(FileUntracked {
                 path: tracked_file_path.to_path_buf(),
             })
@@ -116,7 +125,7 @@ impl RepositoryPaths {
                 history_path,
                 tracked_path: tracked_file_path.to_path_buf(),
             })
-        }
+        })
     }
 }
 
@@ -132,8 +141,10 @@ impl FileDeleted {
             .open(&self.history_path)
     }
 
-    pub fn create_tracked_file(&self, locations: &Locations) -> io::Result<File> {
-        File::create(locations.tracked_from_history(&self.history_path))
+    pub fn create_tracked_file(&self, locations: &Locations) -> Result<File> {
+        Ok(File::create(
+            locations.tracked_from_history(&self.history_path)?,
+        )?)
     }
 }
 
@@ -146,16 +157,16 @@ impl FileUntracked {
         OpenOptions::new().read(true).open(&self.path)
     }
 
-    pub fn create_history_file(&self, locations: &Locations) -> io::Result<File> {
-        let history_path = locations.history_from_tracked(&self.path);
+    pub fn create_history_file(&self, locations: &Locations) -> Result<File> {
+        let history_path = locations.history_from_tracked(&self.path)?;
 
-        history_path.parent().map(|dir_path| {
-            if !dir_path.exists() {
-                fs::create_dir_all(dir_path).unwrap();
+        if let Some(parent_path) = history_path.parent(){
+            if !parent_path.exists() {
+                fs::create_dir_all(parent_path)?;
             }
-        });
+        } 
 
-        File::create(history_path)
+        Ok(File::create(history_path)?)
     }
 }