A virtualenv indicator that works everywhere

NOTE: While none of the significant elements of this approach require zsh, the formatting syntax and mechanism for updating a prompt dynamically differ between zsh and bash. See the end for a version with zsh-specific bits stripped out.

For a while, I’ve been using a variation on this zsh prompt tweak to get a pretty indication that I’m in a virtualenv. However, I was never quite satisfied with it for two reasons:

1. It only works for virtualenvs activated through virtualenvwrapper
2. It goes away if I launch a child shell… which is when I’m most likely to be confused and needing an indicator.

The solution was obvious: Instead of using a virtualenvwrapper hook, put something in my .zshrc which will detect virtualenvs opened through any means.

For those who just want something to copy-paste, here’s what I came up with:

zsh_virtualenv_prompt() {
    # If not in a virtualenv, print nothing
    [[ "$VIRTUAL_ENV" == "" ]] && return

    # Support both ~/.virtualenvs/<name> and <name>/venv
    local venv_name="${VIRTUAL_ENV##*/}"
    if [[ "$venv_name" == "venv" ]]; then
        venv_name=${VIRTUAL_ENV%/*}
        venv_name=${venv_name##*/}
    fi

    # Distinguish between the shell where the virtualenv was activated and its
    # children
    if typeset -f deactivate >/dev/null; then
        echo "[%F{green}${venv_name}%f] "
    else
        echo "<%F{green}${venv_name}%f> "
    fi
}

setopt PROMPT_SUBST PROMPT_PERCENT

# Display a "we are in a virtualenv" indicator that works in child shells too
VIRTUAL_ENV_DISABLE_PROMPT=1
RPS1='$(zsh_virtualenv_prompt)'

First, notice the use of VIRTUAL_ENV_DISABLE_PROMPT. This is because activate will prepend a less attractive indicator to PS1 that also goes away in child shells.

(Just make sure you remove any PS1="$_OLD_VIRTUAL_PS1" you might have added to postactivate or you’ll have no prompt after typing workon projname and be very confused.)

Second, note the use of PROMPT_SUBST. This is actually shared with my code for adding git branch information to PS1, PS2, and PS3 because profiling showed it to be faster than using a precmd function.

Third, note the single quotes for RPS1. That’s necessary to defer the invocation of $(check_virtualenv) so PROMPT_SUBST can see it.

I also added a couple of convenience features:

• I have had a history of virtualenvwrapper not getting along with Python 3.x, so some of my projects have their virtualenvs at ~/src/<name>/venv rather than ~/.virtualenvs/<name>. This script will display <name> in the prompt either way.
• If I’m in a child shell where the deactivate function isn’t available, the prompt will show <foo> rather than [foo] to make me aware of that.

Aside from that, it’s just ordinary efforts to avoid performing disk I/O or use $() in something that’s going to get run every time the prompt is displayed, and a function structured so the most common code path executes the fewest statements.

While this StackOverflow answer cautions against using VIRTUAL_ENV to detect virtualenvs, its reasoning doesn’t apply here, because it’s talking about detecting whether your Python script is running under the influence of a virtualenv, regardless of whether activate was used to achieve that. The purpose of this indicator, on the other hand, is specifically to detect the effects of activate so I don’t run something like manage.py runserver or pip install in the wrong context.

Bash Version

bash_virtualenv_prompt() {
    # If not in a virtualenv, print nothing
    [[ "$VIRTUAL_ENV" == "" ]] && return

    # Support both ~/.virtualenvs/<name> and <name>/venv
    local venv_name="${VIRTUAL_ENV##*/}"
    if [[ "$venv_name" == "venv" ]]; then
        venv_name=${VIRTUAL_ENV%/*}
        venv_name=${venv_name##*/}
    fi

    # Distinguish between the shell where the virtualenv was activated and its
    # children
    if typeset -f deactivate >/dev/null; then
        echo "[${venv_name}] "
    else
        echo "<${venv_name}> "
    fi
}

# Display a "we are in a virtualenv" indicator that works in child shells too
VIRTUAL_ENV_DISABLE_PROMPT=1
PS1='$(bash_virtualenv_prompt)'"$PS1"

It’s almost identical to the zsh version, but the following functions which zsh provides for free are left to the reader in the bash version:

• Implementing a right-aligned chunk of the prompt which stays properly positioned if you resize your terminal.
• Using tput to retrieve the colour-setting escape sequences for your terminal and then caching them in a variable so you’re neither hard-coding for a specific terminal type nor performing multiple subprocess calls each time you display your prompt.

A virtualenv indicator that works everywhere by Stephan Sokolow is licensed under a Creative Commons Attribution-ShareAlike 4.0 International License.