Made the `context` segment customizable.

2 files changed, 13 insertions, 2 deletions

diff --git a/README.md b/README.md
index e38b04b4..53c8f482 100644
--- a/README.md
+++ b/README.md
@@ -366,7 +366,13 @@ of the `time` segment, you would use:
 
 Use the segment names from the above section `Segment Customization`.  Some of
 the Segments have special color variables, as they change the colors according
-to some internal rules. These Segments are `vcs`, `rspec_stats`, `symfony2_tests`:
+to some internal rules. These Segments are `context`, `vcs`, `rspec_stats`, `symfony2_tests`:
+
+    # Colors for root and other users:
+    POWERLEVEL9K_CONTEXT_DEFAULT_BACKGROUND="green"
+    POWERLEVEL9K_CONTEXT_DEFAULT_FOREGROUND="cyan"
+    POWERLEVEL9K_CONTEXT_ROOT_BACKGROUND="red"
+    POWERLEVEL9K_CONTEXT_ROOT_FOREGROUND="blue"
 
     # General VCS color segments:
     POWERLEVEL9K_VCS_FOREGROUND='blue'
diff --git a/powerlevel9k.zsh-theme b/powerlevel9k.zsh-theme
index 4f75e89b..8677ac4a 100644
--- a/powerlevel9k.zsh-theme
+++ b/powerlevel9k.zsh-theme
@@ -437,7 +437,12 @@ prompt_aws() {
 # Note that if $DEFAULT_USER is not set, this prompt segment will always print
 prompt_context() {
   if [[ "$USER" != "$DEFAULT_USER" || -n "$SSH_CLIENT" ]]; then
-    $1_prompt_segment "$0" "$DEFAULT_COLOR" "011" "%(!.%{%F{yellow}%}.)$USER@%m"
+    if [[ $(print -P "%#") == '#' ]]; then
+      # Shell runs as root
+      $1_prompt_segment "$0_ROOT" "$DEFAULT_COLOR" "yellow" "$USER@%m"
+    else
+      $1_prompt_segment "$0_DEFAULT" "$DEFAULT_COLOR" "011" "$USER@%m"
+    fi
   fi
 }