zsh and virtualenv

A week ago or so I finally got off my arse and did the pragmatic programmer thing, setting aside those measly ten minutes to check out virtualenv (well, I also checked out buildout, but I won’t discuss it in this post). I knew pretty much what to expect, but I wanted to get my hands dirty with them so I could see what I assumed I’ve been missing out on for so long (and indeed I promptly kicked myself for not doing it sooner, yada yada, you probably know the drill about well-known-must-know-techniques-and-tools-that-somehow-you-don’t-know). Much as I liked virtualenv, there were two things I didn’t like about environment activation in virtualenv. First, I found typing ‘source bin/activate’ (or similar) cumbersome, I wanted something short and snazzy that worked regardless of where inside the virtualenv I am so long as I’m somewhere in it (it makes sense to me to say that I’m ‘in’ a virtualenv when my current working directory is somewhere under the virtualenv’s directory). Note that being “in” a virtualenv isn’t the same as activating it; you can change directory from virtualenv foo to virtualenv bar, and virtualenv foo will remain active. Indeed, this was the second problem I had: I kept forgetting to activate my virtualenv as I started using it or to deactivate the old one as I switched from one to another.

zsh to the rescue. You may recall that I already mentioned the tinkering I’ve done to make it easier to remember my current DVCS branch. Briefly, I have a function called _rprompt_dvcs which is evaluated whenever zsh displays my prompt and if I’m in a git/Mercurial repository it sets my right prompt to the name of the current branch in blue/green. You may also recall that while I use git itself to tell me if I’m in a git repository at all and which branch I’m at (using git branch --no-color 2> /dev/null | sed -e '/^[^*]/d' -e 's/* \(.*\)/(\1)/'), I had to resort to a small C program (fast_hg_root) in order to decide whether I’m in a Mercurial repository or not and then I manually parse the branch with cat. As I said in the previous post about hg and prompt, I’m not into giving hg grief about speed vs. git, but when it comes to the prompt things are different.

With this background in mind, I was perfectly armed to solve my woes with virtualenv. First, I changed fast_hg_root to be slightly more generic and search for a user-specified “magic” filename upwards from the current working directory (I called the outcome walkup, it’s really simple and nothing to write home about…). For example, to mimic fast_hg_root with walkup, you’d run it like so: $ walkup .hg. Using $ walkup bin/activate to find my current virtualenv (if any at all), I could easily add the following function to my zsh environment:
[sourcecode language=”bash”]
act () {
if [ -n "$1" ]
then
if [ ! -d "$1" ]
then
echo "act: $1 no such directory"
return 1
fi
if [ ! -e "$1/bin/activate" ]
then
echo "act: $1 is not a virtualenv"
return 1
fi
if which deactivate > /dev/null
then
deactivate
fi
cd "$1"
source bin/activate
else
virtualenv="$(walkup bin/activate)"
if [ $? -eq 1 ]
then
echo "act: not in a virtualenv"
return 1
fi
source "$virtualenv"/bin/activate
fi
}
[/sourcecode]
Now I can type $ act anywhere I want in a virtualenv, and that virtualenv will become active; this saves figuring out the path to bin/activate and ending up typing something ugly like $ source ../../bin/activate. If you want something that can work for you without a special binary on your host, there’s also a pure-shell version of the same function in the collapsed snippet below.
[sourcecode language=”bash” collapse=”true”]
function act() {
if [ -n "$1" ]; then
if [ ! -d "$1" ]; then
echo "act: $1 no such directory"
return 1
fi
if [ ! -e "$1/bin/activate" ]; then
echo "act: $1 is not a virtualenv"
return 1
fi

if which deactivate > /dev/null; then
deactivate
fi
cd "$1"
source bin/activate
else
stored_dir="$(pwd)"
while [ ! -f bin/activate ]; do
if [ $(pwd) = / ]; then
echo "act: not in a virtualenv"
cd "$stored_dir"
return 1
fi
cd ..
done
source bin/activate
cd "$stored_dir"
fi
}
[/sourcecode]

This was nice, but only solved half the problem: I still kept forgetting to activate the virtualenv, or moving out of a virtualenv and forgetting that I left it activated (this can cause lots of confusion, for example, if you’re simultaneously trying out this, this, this or that django-facebook integration modules, more than one of them thinks that facebook is a good idea for a namespace to take!). To remind me, I wanted my left prompt to reflect my virtualenv in the following manner (much like my right prompt reflects my current git/hg branch if any):

1. If I’m not in a virtualenv and no virtualenv is active, do nothing.
2. If I’m in a virtualenv and it is not active, display its name as part of the prompt in white.
3. If I’m in a virtualenv and it is active, display its name as part of the prompt in green.
4. If I’m not in a virtualenv but some virtualenv is active, display its name in yellow.
5. Finally, if I’m in one virtualenv but another virtualenv is active, display both their names in red.

So, using walkup, I wrote the virtualenv parsing functions:
[sourcecode language=”bash”]
function active_virtualenv() {
if [ -z "$VIRTUAL_ENV" ]; then
# not in a virtualenv
return
fi

basename "$VIRTUAL_ENV"
}

function enclosing_virtualenv() {
if ! which walkup > /dev/null; then
return
fi
virtualenv="$(walkup bin/activate)"
if [ -z "$virtualenv" ]; then
# not in a virtualenv
return
fi

basename $(grep VIRTUAL_ENV= "$virtualenv"/bin/activate | sed -E ‘s/VIRTUAL_ENV="(.*)"$/\1/’)
}
[/sourcecode]

All that remained was to change my lprompt function to look like so (remember I have setopt prompt_subst on):
[sourcecode language=”bash”]
function _lprompt_env {
local active="$(active_virtualenv)"
local enclosing="$(enclosing_virtualenv)"
if [ -z "$active" -a -z "$enclosing" ]; then
# no active virtual env, no enclosing virtualenv, just leave
return
fi
if [ -z "$active" ]; then
local color=white
local text="$enclosing"
else
if [ -z "$enclosing" ]; then
local color=yellow
local text="$active"
elif [ "$enclosing" = "$active" ]; then
local color=green
local text="$active"
else
local color=red
local text="$active":"$enclosing"
fi
fi
local result="%{$fg[$color]%}${text}$rst "
echo -n $result
}

function lprompt {
local col1 col2 ch1 ch2
col1="%{%b$fg[$2]%}"
col2="%{$4$fg[$3]%}"
ch1=$col1${1[1]}
ch2=$col1${1[2]}

local _env=’$(_lprompt_env)’

local col_b col_s
col_b="%{$fg[green]%}"
col_s="%{$fg[red]%}"

PROMPT="\
$bgc$ch1\
$_env\
%{$fg_bold[white]%}%m:\
$bgc$col2%B%1~%b\
$ch2$rst \
$col2%#$rst "
}
[/sourcecode]

A bit lengthy, but not very difficult. I suffered a bit until I figured out that I should escape the result of _lprompt_virtualenv using a percent sign (like so: "%{$fg[$color]%}${text}$rst "), or else the ANSII color escapes are counted for cursor positioning purposes and screw up the prompt’s alignment. Meh. Also, remember to set VIRTUAL_ENV_DISABLE_PROMPT=True somewhere, so virtualenv’s simple/default prompt manipulation functionality won’t kick in and screw things up for you, and we’re good to go.

The result looks like so (I still don’t know how to do a terminal-“screenshot”-to-html, here’s a crummy png):

Voila! Feel free to use these snippets, and happy zshelling!