added ssh support and POWERLEVEL9K_HIDE_HOST feature

refactorized the sourcecode

added true/false switch

2 files changed, 21 insertions, 4 deletions

diff --git a/README.md b/README.md
index dc98b7d8..147aee40 100644
--- a/README.md
+++ b/README.md
@@ -241,6 +241,11 @@ elements (it is by default), and define a `DEFAULT_USER` in your `~/.zshrc`:
 |----------|---------------|-------------|
 |`DEFAULT_USER`|None|Username to consider a "default context" (you can also use `$USER`)|
 
+You can use POWERLEVEL9K_HIDE_HOST for hiding the hostname in the prompt
+when you are not in a ssh session. For example:
+
+    POWERLEVEL9K_HIDE_HOST="yes"
+
 ##### dir
 
 The `dir` segment shows the current working directory. When using the "Awesome
diff --git a/powerlevel9k.zsh-theme b/powerlevel9k.zsh-theme
index c4b24b55..f2262994 100755
--- a/powerlevel9k.zsh-theme
+++ b/powerlevel9k.zsh-theme
@@ -434,13 +434,25 @@ prompt_battery() {
 # Context: user@hostname (who am I and where am I)
 # Note that if $DEFAULT_USER is not set, this prompt segment will always print
 prompt_context() {
-  if [[ "$USER" != "$DEFAULT_USER" || -n "$SSH_CLIENT" ]]; then
+  local current_state="DEFAULT"
+  declare -A context_states
+  context_states=(
+    "ROOT"      "yellow"
+    "DEFAULT"   "011"
+  )
+  local content="$USER"
+  if [[ "$USER" != "$DEFAULT_USER" ]]; then
     if [[ $(print -P "%#") == '#' ]]; then
-      # Shell runs as root
-      "$1_prompt_segment" "$0_ROOT" "$2" "$DEFAULT_COLOR" "yellow" "$USER@%m"
+      current_state="ROOT"
+    fi
+    if [[ -z "$SSH_CLIENT" && -z "$SSH_TTY" ]]; then
+      if [[ "$POWERLEVEL9K_HIDE_HOST" == "false" || -z "$POWERLEVEL9K_HIDE_HOST"  ]]; then
+        content="${content}@%m"
+      fi
     else
-      "$1_prompt_segment" "$0_DEFAULT" "$2" "$DEFAULT_COLOR" "011" "$USER@%m"
+      content="${content}@%m"
     fi
+    "$1_prompt_segment" "${0}_${current_state}" "$2" "$DEFAULT_COLOR" "${context_states[$current_state]}" "${content}"
   fi
 }