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;

// Configuration.

// What to print before and after the current branch name.
const BRANCH_PREFIX: &str = "\u{e725} ";    // Nerd Font nf-dev-git_branch
const BRANCH_SUFFIX: &str = " ";

// What to print before the number of commits ahead and behind remote.
const AHEAD_PREFIX: &str = "\u{f55c}";      // Nerd Font nf-mdi-arrow_up
const BEHIND_PREFIX: &str = "\u{f544}";     // Nerd Font nf-mdi-arrow_down

// What to print between the current branch information and the file
// status information.
const STATUS_SEPARATOR: &str = " | ";       // U+007C VERTICAL LINE

// What to print before the number of stashes.
const STASHED_PREFIX: &str = "\u{2691}";    // U+2691 BLACK FLAG

// What to print before the number of staged files.
const STAGED_PREFIX: &str = "\u{ea71}";     // Nerd Font nf-cod-circle_filled

// What to print before the number of changed files.
const CHANGED_PREFIX: &str = "+";           // U+002B PLUS SIGN

// What to print before the number of conflicting files.
const CONFLICTED_PREFIX: &str = "\u{f655}"; // Nerd font nf-mdi-close

// What to print before the number of untracked files.
const UNTRACKED_PREFIX: &str = "\u{f142}";  // Nerd Font nf-fa-ellipsis_v

// What to print before the number of ignored files.
const IGNORED_PREFIX: &str = "!";           // U+0021 EXCLAMATION MARK

// What to print if the branch is clean.
const CLEAN_INDICATOR: &str = "\u{f00c}";   // Nerd Font nf-fa-check

// Styling.

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
}