Rewrite the prompt to accept arguments

src/main.rs

@@ -3,22 +3,37 @@ extern crate clap;
 extern crate git2;
 extern crate regex;
 
-use clap::{App, Arg};
+use clap::{App, AppSettings, Arg, SubCommand};
 
 mod prompt;
 mod precmd;
 
 fn main() {
     let matches = App::new("Pure prompt")
-        .arg(Arg::with_name("part")
-            .index(1)
-            .possible_values(&["precmd", "prompt"])
-            .required(true))
-        .get_matches();
-
-    match matches.value_of("part").unwrap() {
-        "precmd" => precmd::display(),
-        "prompt" => prompt::display(),
-        _ => (),
+      .setting(AppSettings::SubcommandRequired)
+      .subcommand(SubCommand::with_name("precmd"))
+      .subcommand(
+        SubCommand::with_name("prompt")
+          .arg(
+            Arg::with_name("last_return_code")
+              .short("r")
+              .takes_value(true)
+          )
+          .arg(
+            Arg::with_name("keymap")
+              .short("k")
+              .takes_value(true)
+          ),
+      )
+      .get_matches();
+
+    match matches.subcommand() {
+      ("precmd", _) => precmd::display(),
+      ("prompt", Some(sub_matches)) => {
+        let last_return_code = sub_matches.value_of("last_return_code").unwrap();
+        let keymap = sub_matches.value_of("keymap").unwrap();
+        prompt::display(last_return_code, keymap);
+      }
+      _ => (),
     }
-}
+}

src/prompt.rs

@@ -1,18 +1,17 @@
-use std::env;
+const INSERT_SYMBOL:&str = "❯";
+const COMMAND_SYMBOL:&str = "⬢";
 
-const SYMBOL:&str = "❯";
-
-pub fn display() {
-  let var = env::var("_LAST_RETURN_STATUS");
-  let last_command = match var {
-    Ok(val) => val.parse::<u8>().unwrap(),
-    Err(_e) => 1
+pub fn display(return_code: &str, keymap: &str) {
+  let symbol = match keymap {
+    "vicmd" => COMMAND_SYMBOL,
+    _ => INSERT_SYMBOL,
   };
 
-  let shell_color = match last_command {
-    0 => 5,
+  let shell_color = match (symbol, return_code) {
+    (COMMAND_SYMBOL, _) => 3,
+    (_, "0") => 5,
     _ => 9,
   };
 
-  print!("%F{{{}}}{}%f", shell_color, SYMBOL);
+  print!("%F{{{}}}{}%f ", shell_color, symbol);
 }