Newer
Older
zsh-rust-git-prompt / src / main.rs
use anes::{Attribute, Color, ResetAttributes, SetAttribute, SetForegroundColor};
use std::io;
use std::process::ExitCode;
use regex::Regex;
use std::num::NonZeroUsize;
use substring::Substring;

const BRANCH_PREFIX: &str = "\u{e725} ";
const BRANCH_SUFFIX: &str = " ";

const AHEAD_PREFIX: &str = "\u{f55c}";
const BEHIND_PREFIX: &str = "\u{f544}";

const STATUS_SEPARATOR: &str = " | ";

const STASHED_PREFIX: &str = "⚑";
const STAGED_PREFIX: &str = "\u{ea71}";
const CHANGED_PREFIX: &str = "+";
const CONFLICTED_PREFIX: &str = "\u{f655}";
const UNTRACKED_PREFIX: &str = "\u{f142}";
const IGNORED_PREFIX: &str = "!";

const CLEAN_INDICATOR: &str = "\u{f00c}";

fn reset_style() -> String {
    return format!("%{{{}%}}", ResetAttributes);
}

fn branch_style() -> String {
    return format!("%{{{}{}{}%}}", SetForegroundColor(Color::Black), SetAttribute(Attribute::Bold), SetAttribute(Attribute::Underline));
}

fn stashed_style() -> String {
    return format!("%{{{}%}}", SetForegroundColor(Color::Blue));
}

fn staged_style() -> String {
    return format!("%{{{}%}}", SetForegroundColor(Color::Blue));
}

fn changed_style() -> String {
    return format!("%{{{}%}}", SetForegroundColor(Color::DarkMagenta));
}

fn conflicted_style() -> String {
    return format!("%{{{}%}}", SetForegroundColor(Color::DarkRed));
}

fn untracked_style() -> String {
    return format!("%{{{}%}}", SetForegroundColor(Color::DarkGreen));
}

fn ignored_style() -> String {
    return format!("%{{{}%}}", SetForegroundColor(Color::Gray));
}

fn clean_style() -> String {
    return format!("%{{{}%}}", SetForegroundColor(Color::DarkGreen));
}

fn main() -> ExitCode {
    let mut git_status = String::new();

    io::stdin().read_line(&mut git_status).expect("Failed to read input");

    if git_status.len() == 0 {
        print!("");
        return ExitCode::SUCCESS;
    }

    // Current commit.
    // example: # branch.oid 825882679bbf52650f8115bfec8591d973d80b49
    let hash = match Regex::new(r"# branch\.oid ([^\x00]+)").unwrap().captures(&git_status) {
        Some(caps) => caps.get(1).unwrap().as_str(),
        None => "???",
    };

    // Current branch.
    // example: # branch.head master
    // example: # branch.head (detached)
    let branch = match Regex::new(r"# branch\.head ([^\x00]+)").unwrap().captures(&git_status) {
        Some(caps) => {
            let name = caps.get(1).unwrap().as_str();
            if name == "(detached)" {
                // use the (shortened) commit hash
                format!("{BRANCH_PREFIX}:{}{}{}{BRANCH_SUFFIX}", branch_style(), hash.substring(0, 7), reset_style())
            } else {
                format!("{BRANCH_PREFIX}{}{}{}{BRANCH_SUFFIX}", branch_style(), name, reset_style())
            }
        },
        None => "???".to_string(),
    };

    // Commits ahead/behind of upstream branch.
    // example: # branch.ab +2 -5
    let (ahead, behind) = match Regex::new(r"# branch\.ab \+(\d+) -(\d+)").unwrap().captures(&git_status) {
        Some(caps) => (
            format!("{AHEAD_PREFIX}{}", caps.get(1).unwrap().as_str()),
            format!("{BEHIND_PREFIX}{}", caps.get(2).unwrap().as_str())
        ),
        None => (
            format!("{AHEAD_PREFIX}0"),
            format!("{BEHIND_PREFIX}0")
        ),
    };
    
    // Number of stashed files.
    // example: # stash 4
    let stashed = match Regex::new(r"# stash (\d+)").unwrap().captures(&git_status) {
        Some(caps) => format!("{}{STASHED_PREFIX}{}{} ", stashed_style(), caps.get(1).unwrap().as_str(), reset_style()),
        None => "".to_string(),
    };

    // Number of staged files.
    // example: 1 A. N... 000000 100644 100644 0000000000000000000000000000000000000000 e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 file.staged
    let staged = match NonZeroUsize::new(Regex::new(r"1 A. ([^\x00]+)").unwrap().find_iter(&git_status).count()) {
        Some(n) => format!("{}{STAGED_PREFIX}{n}{} ", staged_style(), reset_style()),
        None => "".to_string(),
    };

    // Number of changed files.
    // example: 1 .M N... 000000 100644 100644 0000000000000000000000000000000000000000 e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 file.changed
    let changed = match NonZeroUsize::new(Regex::new(r"1 .M ([^\x00]+)").unwrap().find_iter(&git_status).count()) {
        Some(n) => format!("{}{CHANGED_PREFIX}{n}{} ", changed_style(), reset_style()),
        None => "".to_string(),
    };

    // Number of conflicted files.
    // example: u UD N... 000000 100644 100644 0000000000000000000000000000000000000000 e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 file.conflicted
    let conflicted = match NonZeroUsize::new(Regex::new(r"u .. ([^\x00]+)").unwrap().find_iter(&git_status).count()) {
        Some(n) => format!("{}{CONFLICTED_PREFIX}{n}{} ", conflicted_style(), reset_style()),
        None => "".to_string(),
    };
    
    // Number of untracked files.
    // example: ? file.untracked
    let untracked = match NonZeroUsize::new(Regex::new(r"\? ([^\x00]+)").unwrap().find_iter(&git_status).count()) {
        Some(n) => format!("{}{UNTRACKED_PREFIX}{n}{} ", untracked_style(), reset_style()),
        None => "".to_string(),
    };

    // Number of ignored files.
    // example: ! file.ignored
    let ignored = match NonZeroUsize::new(Regex::new(r"! ([^\x00]+)").unwrap().find_iter(&git_status).count()) {
        Some(n) => format!("{}{IGNORED_PREFIX}{n}{} ", ignored_style(), reset_style()),
        None => "".to_string(),
    };

    let status_string = format!("{stashed}{ignored}{staged}{changed}{conflicted}{untracked}");

    print!("{branch}{ahead}{behind}{STATUS_SEPARATOR}{}", 
        if status_string.len() > 0 {
            status_string.trim().to_string()
        } else {
            format!("{}{CLEAN_INDICATOR}{}", clean_style(), reset_style())
        }
    );

    ExitCode::SUCCESS
}