From c2ef915921c6b63dc09f39ededc5145bfd2d3009 Mon Sep 17 00:00:00 2001
From: HimDek <himprakashd@gmail.com>
Date: Sun, 29 May 2022 08:28:15 +0530
Subject: [PATCH] Improved mercurial prompt

---
 plugins/mercurial/README.md            | 16 +++--
 plugins/mercurial/mercurial.plugin.zsh | 98 ++++++++++++++++----------
 2 files changed, 71 insertions(+), 43 deletions(-)

diff --git a/plugins/mercurial/README.md b/plugins/mercurial/README.md
index 19d4a17b1..21ad3bba6 100644
--- a/plugins/mercurial/README.md
+++ b/plugins/mercurial/README.md
@@ -52,10 +52,18 @@ plugins=(... mercurial)
 You can also redefine additional vars used in the plugin (after Oh My Zsh is sourced):
 
 ```zsh
-ZSH_THEME_HG_PROMPT_PREFIX="%{$fg_bold[magenta]%}hg:(%{$fg[red]%}"
-ZSH_THEME_HG_PROMPT_SUFFIX="%{$reset_color%}"
-ZSH_THEME_HG_PROMPT_DIRTY="%{$fg[magenta]%}) %{$fg[yellow]%}✗%{$reset_color%}"
-ZSH_THEME_HG_PROMPT_CLEAN="%{$fg[magenta]%})"
+ZSH_THEME_HG_PROMPT_PREFIX="("
+ZSH_THEME_HG_PROMPT_SUFFIX=")"
+ZSH_THEME_HG_PROMPT_SEPARATOR="|"
+ZSH_THEME_HG_PROMPT_BRANCH="%{$fg_bold[magenta]%}"
+ZSH_THEME_HG_PROMPT_BEHIND="%{↓%G%}"
+ZSH_THEME_HG_PROMPT_AHEAD="%{↑%G%}"
+ZSH_THEME_HG_PROMPT_MODIFIED="%{$fg[red]%}%{●%G%}"
+ZSH_THEME_HG_PROMPT_ADDED="%{$fg[blue]%}%{✚%G%}"
+ZSH_THEME_HG_PROMPT_REMOVED="%{$fg[red]%}%{✖%G%}"
+ZSH_THEME_HG_PROMPT_DELETED="%{$fg[red]%}%{🗑️%G%}"
+ZSH_THEME_HG_PROMPT_UNKNOWN="%{$fg[cyan]%}%{…%G%}"
+ZSH_THEME_HG_PROMPT_CLEAN="%{$fg_bold[green]%}%{✔%G%}"
 ```
 
 ### Display repo branch and directory status in prompt
diff --git a/plugins/mercurial/mercurial.plugin.zsh b/plugins/mercurial/mercurial.plugin.zsh
index 9b85d02f6..445062e71 100644
--- a/plugins/mercurial/mercurial.plugin.zsh
+++ b/plugins/mercurial/mercurial.plugin.zsh
@@ -73,46 +73,66 @@ function hg_get_bookmark_name() {
   echo "$(<"$dir/.hg/bookmarks.current")"
 }
 
+function update_hg_vars() {
+  if $(hg id >/dev/null 2>&1); then
+    local rev="$(hg id -n 2>/dev/null | sed 's/[^-0-9]//g')"
+    local branch="$(hg id -b 2>/dev/null)"
+    local hg_status=`hg st`
+    HG_REV_BRANCH="${rev:gs/%/%%}%{$FG[239]%}@%{$reset_color%}${branch:gs/%/%%}"
+    HG_UNKNOWN="$(echo $hg_status | grep "^\?" | wc -l)"
+    HG_MODIFIED="$(echo $hg_status | grep "^\M" | wc -l)"
+    HG_ADDED="$(echo $hg_status | grep "^\A" | wc -l)"
+    HG_REMOVED="$(echo $hg_status | grep "^\R" | wc -l)"
+    HG_DELETED="$(echo $hg_status | grep "^\!" | wc -l)"
+    HG_AHEAD="$(hg log -r "draft()" | grep "summary" | wc -l)"
+    # HG_BEHIND="$(hg incoming | grep "summary" | wc -l)"
+  fi
+}
+
 function hg_prompt_info {
-  local dir branch dirty
-  if ! dir=$(hg_root); then
-    return
+  update_hg_vars
+  if $(hg id >/dev/null 2>&1); then
+    local STATUS="$ZSH_THEME_HG_PROMPT_PREFIX$ZSH_THEME_HG_PROMPT_BRANCH$HG_REV_BRANCH%{${reset_color}%}"
+    if [[ "$HG_BEHIND" -ne "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_BEHIND$HG_BEHIND%{$reset_color%}"
+    fi
+    if [[ "$HG_AHEAD" -ne "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_AHEAD$HG_AHEAD%{$reset_color%}"
+    fi
+    STATUS="$STATUS$ZSH_THEME_HG_PROMPT_SEPARATOR"
+    if [[ "$HG_MODIFIED" -ne "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_MODIFIED$HG_MODIFIED%{$reset_color%}"
+    fi
+    if [[ "$HG_ADDED" -ne "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_ADDED$HG_ADDED%{$reset_color%}"
+    fi
+    if [[ "$HG_REMOVED" -ne "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_REMOVED$HG_REMOVED%{$reset_color%}"
+    fi
+    if [[ "$HG_DELETED" -ne "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_DELETED$HG_DELETED%{$reset_color%}"
+    fi
+    if [[ "$HG_UNKNOWN" -ne "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_UNKNOWN$HG_UNKNOWN%{$reset_color%}"
+    fi
+    if [[ "$HG_MODIFIED" -eq "0" && "$HG_ADDED" -eq "0" && "$HG_REMOVED" -eq "0" && "$HG_DELETED" -eq "0" && "$HG_UNKNOWN" -eq "0" ]]; then
+      STATUS="$STATUS$ZSH_THEME_HG_PROMPT_CLEAN"
+    fi
+    STATUS="$STATUS%{${reset_color}%}$ZSH_THEME_HG_PROMPT_SUFFIX"
+    echo "$STATUS"
   fi
-
-  if [[ ! -f "$dir/.hg/branch" ]]; then
-    branch=default
-  else
-    branch="$(<"$dir/.hg/branch")"
-  fi
-
-  dirty="$(hg_dirty)"
-
-  echo "${ZSH_THEME_HG_PROMPT_PREFIX}${branch:gs/%/%%}${dirty}${ZSH_THEME_HG_PROMPT_SUFFIX}"
 }
 
-function hg_dirty {
-  # Do nothing if clean / dirty settings aren't defined
-  if [[ -z "$ZSH_THEME_HG_PROMPT_DIRTY" && -z "$ZSH_THEME_HG_PROMPT_CLEAN" ]]; then
-    return
-  fi
-
-  # Check if there are modifications
-  local hg_status
-  if [[ "$DISABLE_UNTRACKED_FILES_DIRTY" = true ]]; then
-    if ! hg_status="$(hg status -q 2>/dev/null)"; then
-      return
-    fi
-  else
-    if ! hg_status="$(hg status 2>/dev/null)"; then
-      return
-    fi
-  fi
-
-  # grep exits with 0 when dirty
-  if command grep -Eq '^\s*[ACDIMR!?L].*$' <<< "$hg_status"; then
-    echo $ZSH_THEME_HG_PROMPT_DIRTY
-    return
-  fi
-
-  echo $ZSH_THEME_HG_PROMPT_CLEAN
-}
+# Default values for the appearance of the prompt.
+ZSH_THEME_HG_PROMPT_PREFIX="("
+ZSH_THEME_HG_PROMPT_SUFFIX=")"
+ZSH_THEME_HG_PROMPT_SEPARATOR="|"
+ZSH_THEME_HG_PROMPT_BRANCH="%{$fg_bold[magenta]%}"
+ZSH_THEME_HG_PROMPT_BEHIND="%{↓%G%}"
+ZSH_THEME_HG_PROMPT_AHEAD="%{↑%G%}"
+ZSH_THEME_HG_PROMPT_MODIFIED="%{$fg[red]%}%{●%G%}"
+ZSH_THEME_HG_PROMPT_ADDED="%{$fg[blue]%}%{✚%G%}"
+ZSH_THEME_HG_PROMPT_REMOVED="%{$fg[red]%}%{✖%G%}"
+ZSH_THEME_HG_PROMPT_DELETED="%{$fg[red]%}%{🗑️%G%}"
+ZSH_THEME_HG_PROMPT_UNKNOWN="%{$fg[cyan]%}%{…%G%}"
+ZSH_THEME_HG_PROMPT_CLEAN="%{$fg_bold[green]%}%{✔%G%}"