zshrc

#===============================================================================
#
#          FILE: .zsh_prompt
#
#         USAGE: Sourced automatically in .zshrc
#
#   DESCRIPTION: In my opinion which I respect very much this may be one of the
#                best zsh prompts on the planet.
#
#===============================================================================

autoload -U colors && colors && autoload -U promptinit && autoload -Uz vcs_info

# Make the colors easier to reference
typeset -AHg FG BG

for color in {000..255}; do
    FG[$color]="%{[38;5;${color}m%}"
    BG[$color]="%{[48;5;${color}m%}"
done

local reset white gray lightblue lightgray brightcyan green red yellow blue cyan lines_color git_repo working_dir pyenv_ver
reset="%{${reset_color}%}"
gray="%{$FG[239]%}"
green="%{$FG[120]%}"
red="%{$FG[196]%}"
yellow="%{$FG[011]%}"
cyan="%{$FG[066]%}"
blue="%{$FG[033]%}"
white="%{$FG[015]%}"
orange=%{$FG[208]%}
violet=%{$FG[135]%}
lightblue="%{$FG[039]%}"
lightgray="%{$FG[007]%}"
brightcyan="%{$FG[087]%}"
brightmagenta=%{$FG[198]%}
brightgreen=%{$FG[082]%}
testcolor=%{$FG[123]%}

#Global color scheme - Change these to the variables above to control your dotfiles environment colors.
lines_color="$brightcyan"
git_repo="$green"
working_dir="$lightblue"
pyenv_ver="$violet"
# npm_ver="$brightcyan"

local -A pr_com            # Associative array
local -a prompt_left_lines # Array parameters

zstyle ":pr_jrock:" mode full
zstyle ':pr_jrock:*' hooks pwd usr vcs venv npm jobs prompt
zstyle ':pr_jrock:*' pwd "%~"

# Set vcs_info options
zstyle ':vcs_info:*' enable git                                           # We are only concerned with git VCS, so enable it
zstyle ':vcs_info:(git*):*' check-for-changes true                        # We want to monitor changes to the repository's
# Format of what we will display for the git repo information.
# %s - The VCS in use (git, hg, svn, etc.).
# %i - The current revision number or identifier. (SHA we only display 10 chars)
# %c - The string from the stagedstr style if there are staged changes in the repository.
# %u - The number of unapplied patches (unapplied-string).
# %b - Information about the current branch.
# %m - A "misc" replacement. It is at the discretion of the backend to decide what this replacement expands to.
#      It is currently used by the hg and git backends to display patch information from the mq and stgit extensions.
# Format of what we will display during a special action on the repo (Ex. Interactive rebase or merge conflict)
zstyle ':vcs_info:(git*)' actionformats "(%b|${red}%a${lines_color}%m)"
zstyle ':vcs_info:(git*)' formats "(%b%m)"
zstyle ':vcs_info:git*+set-message:*' hooks git-statuses git-st

# Run all the prompt hook functions
# (stolen, wholesale, from the excellent hook system in vcs_info)
function pr_run_hooks() {
    local hook func
    local -a hooks

    zstyle -g hooks ":pr_jrock:*" hooks

    (( ${#hooks} == 0 )) && return 0

    for hook in ${hooks} ; do
        func="+pr-${hook}"
        if (( ${+functions[$func]} == 0 )); then
            continue
        fi
        true
        ${func} "$@"
        case $? in
            (0)
                ;;
            (*)
                break
                ;;
        esac
    done
}

# This is our prompt, this is a compressed example it will expand
# to the width of the console
# ┌──(~/src/dotfiles)v(dotfiles)──────(✔)─
# ├──(master [origin/master ] Unstaged)
# └───>
function +pr-mode-full() {
    local i info_line_width return_status_width filler venv_shim npm_shim

    infoline=( ${pr_com[pwd]} ${pr_com[usr]} )

    # # If we are in a directory that has a node_modules subdirectory we want to
    # # display that on the info line
    #  [[ -n ${pr_com[npm]} ]] && infoline[1]=(
    #      ${infoline[1]}
    #      "${lines_color}(${pr_com[npm]}${lines_color})${reset}"
    #  )

    # # The prompt was not taking up the full width of the terminal when displaying npm info
    # # So if we are in a npm directory then we add a one character shim to the width of the filler
    #  npm_shim=0
    #  if [[ -n ${pr_com[npm]} ]]; then
    #      npm_shim+=1
    #  fi

    # If we are in a virtualenv we want to display that on the info line
    [[ -n ${pr_com[venv]} ]] && infoline[1]=(
        ${infoline[1]}
        "${lines_color}(${pr_com[venv]}${lines_color})${reset}"
    )

    # The prompt was not taking up the full width of the terminal when in a virtualenv
    # If we are in a virtualenv then we add a one character shim to the width of the filler
    venv_shim=0
    if [[ -n ${pr_com[venv]} ]]; then
        venv_shim+=1 # Shim the filler to take up the entire width of the prompt
    fi

    # Full-width filler; search/replace color wraps to find real text width
    info_line_width=${(S)infoline//\%\{*\%\}} # search-and-replace color escapes
    info_line_width=${#${(%)info_line_width}} # expand all escapes and count the chars
    return_status_width=3 # We will take up three spaces to display the return code (✔)

    # Set the text string that will be used to fill the width of the terminal filler
    # filler="${lines_color}${(l:$(( $COLUMNS - $info_line_width - $return_status_width + $venv_shim + $npm_shim))::─:)}${reset}"
    filler="${lines_color}${(l:$(( $COLUMNS - $info_line_width - $return_status_width + $venv_shim))::─:)}${reset}"
    infoline[-1]=( ${filler} ${infoline[-1]} )

    # --------------------------
    # Assemble the prompt lines
    # --------------------------
    # Default we will always have a info line and a prompt line
    lines=(
        ${(j::)infoline}
        ${pr_com[prompt]}
    )

    # If we are in a git repo we will have three lines info, git_info, prompt
    [[ -n ${pr_com[vcs]} ]] && lines[1]=(
        ${lines[1]}
        "${lines_color}${pr_com[vcs]}${reset}"

    )

    # Add some connecting lines to the beginning of our prompts
    if [[ -n ${pr_com[vcs]} ]]; then
        lines[1]="${lines_color}┌──${lines[1]}${lines_color}─${reset}"
        lines[2]="${lines_color}├──${lines[2]}${reset}"
        lines[3]="${lines_color}└──➤${lines[3]}${reset}"
    else
        lines[1]="${lines_color}┌──${lines[1]}${lines_color}─${reset}"
        lines[2]="${lines_color}└──➤${lines[2]}${reset}"
    fi

    # And we set the value for the prompt_left_lines
    prompt_left_lines=( ${lines[@]} )
}

# Show info collected from vcs_info
function +pr-vcs() {
    local -a v_vcs

    [[ -n ${vcs_info_msg_0_} ]] && v_vcs=(
        ${git_repo}
        ${vcs_info_msg_0_}
        ${reset}
    )

    pr_com[vcs]=${(j::)v_vcs}
}

# Show virtualenv information
function +pr-venv() {
    local -a v_venv
    python_version=$(python --version 2>&1)
    pr_com[python_version]="${python_version}"

    if [[ -n ${VIRTUAL_ENV} ]]; then
        v_venv=(
            ${pyenv_ver}
            "${pr_com[python_version]}"
            ${reset}
        )
    else
        v_venv=()
    fi

    pr_com[venv]=${(j::)v_venv}
}

# # Show npm information
# function +pr-npm() {
#     local -a v_npm
#     npm_version=$(npm --version 2>&1)
#     if [[ ${NODE_NAME}  != "" ]]; then
#         v_npm=(
#             ${npm_ver}
#             ${NODE_NAME}
#             ${reset}
#         )
#     fi

#     pr_com[npm]=${(j::)v_npm}
# }

# Show number of background jobs, or hide if none
function +pr-jobs() {
    local -a v_jobs
    v_jobs=( "%(1j.${lines_color}%j${reset}.)" )
    pr_com[jobs]=${(j::)v_jobs}
}

#AWS Granted Status
_get_aws_account() {
  aws_account_name="${AWS_PROFILE:u}"
  if [[ -z "$aws_account_name" ]]; then
    echo "${lines_color}─${white}[NO-AUTH]${lines_color}─%{$reset%}"
  elif [[ "$aws_account_name" == *"TEST"* ]]; then
    echo "${lines_color}─${green}[$aws_account_name]${lines_color}─%{$reset%}"
  elif [[ "$aws_account_name" == *"QA"* ]]; then
    echo "${lines_color}─${yellow}[$aws_account_name]${lines_color}─%{$reset%}"
  elif [[ "$aws_account_name" == *"PROD"* ]]; then
    echo "${lines_color}─${red%}[$aws_account_name]${lines_color}─%{$reset%}"
  else
    echo "${lines_color}─${orange}[$aws_account_name]${lines_color}─%{$reset%}"
  fi
}

# Get pulumi backend account number
_get_pulumi_backend_account(){
  if [[ ! -z "$PULUMI_BACKEND_URL" ]]; then
    PULUMI_ACCOUNT_NUMBER=$(echo $PULUMI_BACKEND_URL | cut -d '-' -f4)
    echo $(jq -r ".accountList[] | select( .accountId == \"$PULUMI_ACCOUNT_NUMBER\" ) | .accountName" ~/.aws/sso_accounts.json | tr a-z A-Z)
  fi
}

# Check of Pulumi backend matches current AWS profile
_compare_aws_pulumi_account(){
  PULUMI=$( _get_pulumi_backend_account)
  AWS_ACCOUNT_NAME=${AWS_PROFILE:u}
  if [ ! -z ${AWS_PROFILE+x} ] && [ -n "$PULUMI" ]; then
    if [[ "$AWS_ACCOUNT_NAME" =~ "$PULUMI" ]]; then
      echo "${white}[PULUMI ${green}✔${white}]${lines_color}─%{$reset%}"
    else
      echo "${white}[PULUMI ${red}✘${white}]${lines_color}─%{$reset%}"
    fi
  fi
}

function +pr-prompt() {
    local -a v_pwd i_pwd
    local -a exit_status i_usr i_host exit_status
    account=$(_get_aws_account)
    pulumi=$(_compare_aws_pulumi_account)

    # Add the print working directory logic
    zstyle -g i_pwd ":pr_jrock:*" pwd
    v_pwd+=( ${lines_color}\( )
    [[ -w $PWD ]] && v_pwd+=( ${working_dir} ) || v_pwd+=( ${yellow} )
    v_pwd+=( ${i_pwd} )
    v_pwd+=( ${lines_color}\)${lightgray} )
    v_pwd+=( ${account} )
    v_pwd+=( ${pulumi} )
    v_pwd+=( ${reset} )
    pr_com[pwd]=${(j::)v_pwd}${reset}


    # Add exit status check or x logic
    exit_status=( ${lines_color}\( )
    exit_status+="%(0?.${green}✔.${red}✘)"
    exit_status+=( ${lines_color}\) )
    exit_status+=( ${reset} )

    pr_com[usr]=${(j::)exit_status}
}

# Show remote ref name and number of commits ahead-of or behind
function +vi-git-st() {
    local ahead behind remote
    local -a gitstatus

    # Are we on a remote-tracking branch?
    remote=${$(git rev-parse --verify ${hook_com[branch]}@{upstream} \
        --symbolic-full-name --abbrev-ref 2>/dev/null)}

    if [[ -n ${remote} ]] ; then
        ahead=$(git rev-list ${hook_com[branch]}@{upstream}..HEAD 2>/dev/null | wc -l | sed -e 's/^[ \t]*//')
        (( $ahead )) && gitstatus+=( "${green}+${ahead}${lines_color}" )

        behind=$(git rev-list HEAD..${hook_com[branch]}@{upstream} 2>/dev/null | wc -l | sed -e 's/^[ \t]*//')
        (( $behind )) && gitstatus+=( "${red}-${behind}${lines_color}" )

        user_data[gitstatus]=${gitstatus}
        hook_com[branch]="${hook_com[branch]} [${remote}${(j:/:)gitstatus}]"
    fi
}

# Show the above/behind upstream counts more tersely for the compact display
function +vi-git-st-compact() {
    [[ -n ${user_data[gitstatus]} ]] \
        && hook_com[misc]="@{u}${(j:/:)user_data[gitstatus]}"
}

function +vi-git-statuses() {
    git status -s >| /tmp/gitstatus.txt
    staged=$( cat /tmp/gitstatus.txt | grep -c "^[MARCD]")
    unstaged=$( cat /tmp/gitstatus.txt | grep -c "^.[MARCD]")
    untracked=$( cat /tmp/gitstatus.txt | grep -c "^\?")
    stashes=$(git stash list 2>/dev/null | wc -l | sed -e 's/^[ \t]*//')

    if [[ ${staged} != 0 ]] ; then
        hook_com[misc]+=" ${green}${staged}${lines_color}"
    fi

    if [[ ${unstaged} != 0 ]] ; then
        hook_com[misc]+=" ${red}${unstaged}${lines_color}"
    fi

    if [[ ${untracked} != 0 ]] ; then
        hook_com[misc]+=" ${yellow}${untracked}${lines_color}"
    fi

    if [[ ${stashes} != 0 ]] ; then
        hook_com[misc]+=" ${lines_color}${stashes}${lines_color}"
    fi
}

# --------------------------
# Finally we execute the above prompt functions
# --------------------------

# To be added to the precmd_* array so it is executed before each prompt
function precmd_prompt {
    local func

    # Clear out old values
    pr_com=()
    prompt_left_lines=()

    # Collect needed data
    vcs_info
    pr_run_hooks

    # Use the above data and build the prompt arrays
    func="+pr-mode-full"
    ${func} "$@"

    # Set the prompts
    PROMPT="${(F)prompt_left_lines} "
}