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.

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

This entry was posted in Geek Stuff. Bookmark the permalink.

Leave a Reply

Your email address will not be published.

By submitting a comment here you grant this site a perpetual license to reproduce your words and name/web site in attribution under the same terms as the associated post.

All comments are moderated. If your comment is generic enough to apply to any post, it will be assumed to be spam. Borderline comments will have their URL field erased before being approved.