1 files changed, 50 insertions, 6 deletions

diff --git a/src/main.rs b/src/main.rs
index d7dc687..0603cba 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,29 +15,44 @@ fn main() {
              .long("minimal")
              .help("use minimal variant")
         )
+        .arg(Arg::with_name("zsh")
+             .short("z")
+             .long("zsh")
+             .help("Use ZSH formatting")
+        )
         .get_matches();
     if matches.is_present("minimal") {
-        println!("{}", pista_minimal());
+        println!("{}", pista_minimal(matches.is_present("zsh")));
     } else {
-        println!("{}", pista());
+        println!("{}", pista(matches.is_present("zsh")));
     }
 }
 
-fn pista() -> String {
+fn pista(zsh: bool) -> String {
     let cwd = cwd::cwd();
     let (branch, status) = vcs::vcs_status().unwrap_or(("".into(), "".into()));
     let venv = venv::get_name();
     let prompt_char = prompt_char::get_char();
-    format!("%{{{cwd} {branch} {status}%}} %{{\n{venv}{pchar}%}} ",
+    if zsh {
+        format!("%{{{cwd} {branch} {status}%}} %{{\n{venv}{pchar}%}} ",
+            cwd=cwd,
+            branch=branch,
+            status=status,
+            venv=venv,
+            pchar=prompt_char
+            )
+    } else {
+        format!("{cwd} {branch} {status}\n{venv}{pchar} ",
             cwd=cwd,
             branch=branch,
             status=status,
             venv=venv,
             pchar=prompt_char
             )
+    }
 }
 
-fn pista_minimal() -> String {
+fn pista_minimal(zsh: bool) -> String {
     let cwd = cwd::cwd();
     let vcs_tuple = vcs::vcs_status();
     let mut vcs_component = String::new();
@@ -48,10 +63,39 @@ fn pista_minimal() -> String {
     }
     let venv = venv::get_name();
     let prompt_char = prompt_char::get_char();
-    format!("{cwd}{vcs}{venv}{pchar} ",
+    if zsh {
+        let fmt = format!("{cwd}{vcs}{venv}{pchar} ",
+            cwd=cwd,
+            vcs=vcs_component,
+            venv=venv,
+            pchar=prompt_char
+        );
+        let mut ret = String::new();
+        let mut color = false;
+        for ch in fmt.chars() {
+            if color {
+                if ch == 'm' { // colors always end with m
+                    ret.push_str("m%}");
+                    color = false;
+                } else {
+                    ret.push(ch)
+                }
+            } else {
+                if ch == 0x1b_u8.into() { // ESC char, always starts colors
+                    ret.push_str(&format!("%{{{esc}", esc=ch));
+                    color = true;
+                } else {
+                    ret.push(ch);
+                }
+            }
+        }
+        ret
+    } else {
+        format!("{cwd}{vcs}{venv}{pchar} ",
             cwd=cwd,
             vcs=vcs_component,
             venv=venv,
             pchar=prompt_char
             )
+    }
 }