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):

  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:

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!

Advertisement

Tagged: , , ,

§ 8 Responses to zsh and virtualenv

  • Martin says:

    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 :)

    • Yaniv Aknin says:

      /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?! :)

  • Dj Gilcrease says:

    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

    • Yaniv Aknin says:

      I didn’t use hg branch because of the 88-fold decrease in overall performance and significant increase in resource usage:

      $ pwd
      /tmp/foo/Python-3.1.2/Lib/email/mime
      $ time hg branch
      default
      hg branch  0.14s user 0.03s system 98% cpu 0.176 total
      $ time hg branch
      default
      hg branch  0.14s user 0.03s system 98% cpu 0.176 total
      $ time hg branch
      default
      hg branch  0.14s user 0.03s system 98% cpu 0.176 total
      $ time cat $(walkup .hg)/.hg/branch 2>/dev/null || echo default   
      cat $(walkup .hg)/.hg/branch 2> /dev/null  0.00s user 0.00s system 77% cpu 0.002 total
      default
      $ time cat $(walkup .hg)/.hg/branch 2>/dev/null || echo default
      cat $(walkup .hg)/.hg/branch 2> /dev/null  0.00s user 0.00s system 82% cpu 0.002 total
      default
      $ time cat $(walkup .hg)/.hg/branch 2>/dev/null || echo default
      cat $(walkup .hg)/.hg/branch 2> /dev/null  0.00s user 0.00s system 82% cpu 0.002 total
      default
      $ 
      

      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.

  • Simeon says:

    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.

  • Seth says:

    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 […]

  • Dan Crosta says:

    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.

Leave a Reply to Martin Cancel reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

What’s this?

You are currently reading zsh and virtualenv at NIL: .to write(1) ~ help:about.

meta

%d bloggers like this: