From 541e2d9cc4b7ce8b276f964a8d1ffdf1e4c75571 Mon Sep 17 00:00:00 2001
From: Hugo Wang <w@mitnk.com>
Date: Thu, 24 Feb 2022 13:27:25 +0800
Subject: [PATCH] fix issue in prompt rendering

---
 src/prompt/main.rs | 29 ++++++++++++++++++++++-------
 src/shell.rs       |  8 +++++++-
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/prompt/main.rs b/src/prompt/main.rs
index f27e9c5..adfa27d 100644
--- a/src/prompt/main.rs
+++ b/src/prompt/main.rs
@@ -18,9 +18,13 @@ fn is_suffix_char(c: char) -> bool {
     c == ']' || c == '}'
 }
 
-fn is_prompt_item_char(c: char) -> bool {
+fn is_prompt_item_char(c: char, token: &str) -> bool {
     let s = c.to_string();
-    libs::re::re_contains(&s, r#"^[a-zA-Z_]$"#)
+    if token.is_empty() {
+        libs::re::re_contains(&s, r#"^[a-zA-Z_]$"#)
+    } else {
+        libs::re::re_contains(&s, r#"^[a-zA-Z0-9_]$"#)
+    }
 }
 
 pub fn get_prompt_string() -> String {
@@ -31,7 +35,7 @@ pub fn get_prompt_string() -> String {
 }
 
 fn apply_prompt_item(sh: &shell::Shell, result: &mut String, token: &str) {
-    if let Ok(x) = env::var(token) {
+    if let Some(x) = sh.get_env(token) {
         result.push_str(&x);
         return;
     }
@@ -103,7 +107,7 @@ pub fn render_prompt(sh: &shell::Shell, ps: &str) -> String {
                     token.push(c);
                 }
                 continue;
-            } else if is_prompt_item_char(c) {
+            } else if is_prompt_item_char(c, &token) {
                 token.push(c);
                 continue;
             } else if token.is_empty() {
@@ -145,16 +149,27 @@ pub fn render_prompt(sh: &shell::Shell, ps: &str) -> String {
 
 #[cfg(test)]
 mod tests {
+    use std::env;
     use super::render_prompt;
     use super::shell::Shell;
 
     #[test]
     fn test_render_prompt() {
         let mut sh = Shell::new();
-        sh.set_env("USER", "mitnk");
 
-        assert_eq!("mitnk$\n", render_prompt(&sh, "$USER$${newline}"));
-        assert_eq!("mitnk$\n", render_prompt(&sh, "$USER$$newline"));
+        env::set_var("MY_ARCH", "x86-64");
+        env::set_var("OS", "linux");
+        assert_eq!("x86-64/linux", render_prompt(&sh, "$MY_ARCH/$OS"));
+
+        env::set_var("FOOBAR7", "user135");
+        assert_eq!("user135$\n", render_prompt(&sh, "$FOOBAR7$${newline}"));
+        assert_eq!("user135$\n", render_prompt(&sh, "$FOOBAR7$$newline"));
+
+        sh.set_env("FOOBAR8", "u138");
+        assert_eq!("u138", render_prompt(&sh, "$FOOBAR8"));
+
+        assert_eq!("$67", render_prompt(&sh, "$67"));
+        assert_eq!("<NOT_EXISTS>", render_prompt(&sh, "$NOT_EXISTS"));
 
         assert_eq!("$", render_prompt(&sh, "$"));
         assert_eq!("$$", render_prompt(&sh, "$$"));
diff --git a/src/shell.rs b/src/shell.rs
index c6fc605..4d96882 100644
--- a/src/shell.rs
+++ b/src/shell.rs
@@ -248,10 +248,16 @@ impl Shell {
         }
     }
 
+    /// get *Shell Variable*, or *ENV Variable*.
     pub fn get_env(&self, name: &str) -> Option<String> {
         match self.envs.get(name) {
             Some(x) => Some(x.to_string()),
-            None => None,
+            None => {
+                match env::var(name) {
+                    Ok(x) => Some(x),
+                    Err(_) => None,
+                }
+            }
         }
     }