zsh and virtualenv
2010/10/14 § 8 Comments
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:
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
}
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.
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
}
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):
- If I’m not in a virtualenv and no virtualenv is active, do nothing.
- If I’m in a virtualenv and it is not active, display its name as part of the prompt in white.
- If I’m in a virtualenv and it is active, display its name as part of the prompt in green.
- If I’m not in a virtualenv but some virtualenv is active, display its name in yellow.
- 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:
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/')
}
All that remained was to change my lprompt function to look like so (remember I have setopt prompt_subst on):
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 "
}
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!
NIL: .to write(1) ~ help:about 
So… you are telling us that you just got your hands dirty on virtualenv but you did NOT discover virtualenvwrapper?!?
You should google before you implement something :)
/me feels sheepish
I found virtualenv-commands, somehow missed virtualenvwrappers. While Indeed it looks like wrappers could have saved me some of the work, I’m not sure it would have replaced all or even most of it.
In particular, the current do-stuff-on-entering-directory hacks you can find in wrappers’ documentation refer to checking for a virtualenv immediately in the directory you’re at (look for ./.venv), whereas I want to be anywhere in the tree (recursively look upwards for .venv). Also, most of my work was about zsh coloring, and while it’s mentioned in wrappers’ FAQ as an also-can-do thing, it’s more a suggestion of the idea than something complete and I think I’ll have to write most of the scripts I wrote anyway.
The main thing that ends up being redundant is my act function, but then again, that was vy far the easiest bit to write, since I already had fast_hg_root (the function practically worked the first time I typed it).
All in all, of course I should have bloody looked before I implemented something (and more importantly, before blogging about it…); but where’s the fun in that?! :)
Wondering why the “hg branch” command didn’t work for you to tell you if you are in a hg repository and what branch is active if you are. Other then that very useful I find myself forgetting what virtualenv is active all the time, this should help
I didn’t use hg branch because of the 88-fold decrease in overall performance and significant increase in resource usage:
It may not be much if you type the command now and then to see which branch you’re on, but for every prompt it ends up being noticeably slow, and if you’re also compiling something or otherwise chewing your processor/IO – it’s much slower.
If you’re going to this level to script management of your virtualenvs you should consider virtualenvwrapper. It gives you a (tab completion enabled) workon command to choose which virtual env to activate and also has the convenient commands to cd/ls the site-packages dir of the currently activated virtual environment. It also has convenience commands for creating/deleting virtual envs and adding globally installed packages. Finally there’s a hooks system for scripting entry and exit of a virtualenv. A lot there and I saw something about zsh support in 1.6 – I’ve never used zsh so ymmv.
Nice work!
Zsh comes bundled with functionality for putting VCS information in your prompt called VCS_Info. It’s documented in the zshcontrib manpage or you can look at the code here:
http://zsh.git.sourceforge.net/git/gitweb.cgi?p=zsh/zsh;a=tree;f=Functions/VCS_Info
It uses a mechanism very similar to your pure-shell version of walkup to detect when you’re inside a version-controlled directory. :)
You may find reading through the code for the Mercurial backend interesting; I recently did some work on it and I went to great lengths to keep it as fast as possible. If you opt to invoke Mercurial, it does so only once per prompt or you can configure it to use hexdump instead which is very fast.
The version in the current release of zsh is missing a ton of the cool new features, but you can easily use the new VCS_Info now just by copying the VCS_Info directory and putting it on your fpath. There are instructions at the bottom of this file:
http://zsh.git.sourceforge.net/git/gitweb.cgi?p=zsh/zsh;a=blob_plain;f=Misc/vcs_info-examples
Anyway, thanks for an interesting post. Now I want to add auto-detecting virtualenvs into my own prompt. :D
[…] I have been reading some great posts about customizations of ZSH prompt. In the screenshot below I show how my customized ZSH […]
I’ve recently spent a bit of time customizing virtualenvwrapper through its hooks, and just documented it on my blog: http://late.am/post/2011/12/22/customize-virtualenvwrapper-for-fun-and-profit as a sort of reply to your post.