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!

Advertisements

Switching to mercurial: taming zsh

2010/05/14 § 12 Comments

A quick one, so you’ll know not all my posts must have so many words. PEP 385 is materializing, and its time to learn Mercurial. I can’t say I’m a Mercurial expert, but I thought migrating all my git-oriented-zsh-gizmos would help me along the way. The conversion is almost done and had just one somewhat noteworthy tidbit.

A while back I copied from a friend a rather elaborate zsh prompt (not as elaborate as some people’s…), which includes my current git branch (if any) in it. The code to make it looked like this:

parse_git_branch() {
        git branch --no-color 2> /dev/null | sed -e '/^[^*]/d' -e 's/* \(.*\)/(\1)/'
}

function _rprompt_git {
    local git_branch="$(parse_git_branch)"
    if [ -n "$git_branch" ]; then
        result=":%{$fg[blue]%}$git_branch$rst"
    fi
    echo -n $result
}

I’ve na├»vely added a Mercurial equivalent to parse_git_branch using hg bookmarks or hg branch, and retrofitted/redubbed _rprompt_git into _rprompt_dvcs. It worked well, but was slow. I’m not the kind of person to give hg grief over reasonable speed differences with the speed monster, but you can’t wait 0.15 seconds for your friggin’ prompt, now can you (this is not an invitation for a git/hg/bzr performance holy war in the comments, people). Removing a call to hg by using $ hg root just once, storing that value and using cat to get the actual branch/bookmark didn’t speed things up enough. #mercurial on freenode was kind but didn’t know how to help, other than suggest I buy a really fast computer… Blah, I’ll have to jot something out in C.

One bitbucket account later and a bit of tinkering led to fast_hg_root, which is a dumb C program which acts as a (fast) replacement to hg root. So now my dvcs-prompt-related code includes:

parse_git_branch() {
        git branch --no-color 2> /dev/null | sed -e '/^[^*]/d' -e 's/* \(.*\)/(\1)/'
}

parse_hg_branch() {
        if ! HG_ROOT=$(fast_hg_root) 2> /dev/null; then
            # not an HG repository, quit
            return
        fi
        BOOKMARK=$(cat "$HG_ROOT"/.hg/bookmarks.current 2> /dev/null)
        if [ -n "$BOOKMARK" ]; then
            # have a current bookmark, display that
            echo $BOOKMARK
            return
        fi
        # display the current branch or 'default'
        cat "$HG_ROOT"/.hg/branch 2> /dev/null || echo 'default'
}

function _rprompt_dvcs {
    local git_branch="$(parse_git_branch)"
    if [ -n "$git_branch" ]; then
        result=":%{$fg[blue]%}$git_branch$rst"
    else
        local hg_branch="$(parse_hg_branch)"
        if [ -n "$hg_branch" ]; then
            result=":%{$fg[green]%}$hg_branch$rst"
        fi
    fi
    echo -n $result
}

Which leads to my terminal sessions on my netbook monica looking like so:
DVCS aware zsh prompts

C’est tout. If someone is interested, I’ll open a repository with more of my .zsh.d stuff. By the way, anyone knows of a good way to capture a terminal into HTML, preserving ANSI color? I know of ttyrec and found HTML::FromANSI, but was hoping for a finished program or at least, uhm, a library less in Perl. :p

Where Am I?

You are currently browsing entries tagged with prompt at NIL: .to write(1) ~ help:about.