improve output in non existent directories

2 files changed, 14 insertions, 6 deletions

diff --git a/src/cwd.rs b/src/cwd.rs
index 95b6659..157858c 100644
--- a/src/cwd.rs
+++ b/src/cwd.rs
@@ -2,8 +2,9 @@ use std::env;
 use tico::tico;
 use colored::*;
 
-pub fn cwd() -> colored::ColoredString {
-    let mut path = format!("{}", env::current_dir().unwrap_or("".into()).display());
+pub fn cwd() -> Option<colored::ColoredString> {
+    let path_env = env::current_dir().ok()?;
+    let mut path = format!("{}", path_env.display());
     let home = env::var("HOME").unwrap();
     let tilde_expand = env::var("EXPAND_TILDE").unwrap_or("0".into());
 
@@ -19,8 +20,8 @@ pub fn cwd() -> colored::ColoredString {
     let cwd_shorten = env::var("SHORTEN_CWD").unwrap_or("1".into());
     let cwd_color = env::var("CWD_COLOR").unwrap_or("white".into());
     match cwd_shorten.as_ref() {
-        "0" => return path.color(cwd_color),
-        _ => return tico(&path[..]).color(cwd_color)
+        "0" => return Some(path.color(cwd_color)),
+        _ => return Some(tico(&path[..]).color(cwd_color))
     }
 
 }
diff --git a/src/main.rs b/src/main.rs
index 0603cba..3096635 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,6 +4,7 @@ mod vcs;
 mod venv;
 
 use clap::{Arg, App};
+use colored::*;
 
 fn main() {
     let matches = App::new("Pista")
@@ -29,7 +30,10 @@ fn main() {
 }
 
 fn pista(zsh: bool) -> String {
-    let cwd = cwd::cwd();
+    let cwd = match cwd::cwd() {
+        Some(c) => c,
+        None => "[directory does not exist]".color("red")
+    };
     let (branch, status) = vcs::vcs_status().unwrap_or(("".into(), "".into()));
     let venv = venv::get_name();
     let prompt_char = prompt_char::get_char();
@@ -53,7 +57,10 @@ fn pista(zsh: bool) -> String {
 }
 
 fn pista_minimal(zsh: bool) -> String {
-    let cwd = cwd::cwd();
+    let cwd = match cwd::cwd() {
+        Some(c) => c,
+        None => "[directory does not exist]".color("red")
+    };
     let vcs_tuple = vcs::vcs_status();
     let mut vcs_component = String::new();
     if let Some((branch, status)) = vcs_tuple {