Better process for zsh virtualenvwrapper plugin
By: Ken Lai
Tags: Ken Lai, Tech, Zsh, Python

Virtualenvwrapper is one of the must have scripts for building a Python development environment. With the power of zsh oh-my-zsh plugin framework, we can automate some processes such as activate & deactivate.

But the build-in virtualenvwrapper fails to work sometimes what it sould be. I decided to fix it by myself after googling for existed solutions.

virtualenvwrapper plugin analyze & optimize

The existed virtualenvwrapper plugin is quite useful in most of time, but it would malfunction occasionally. So I want to break it down in this section to find proper ways to improve it.

As it is a plugin for zsh environment, the code is quite straight.

The first part is used to check system variables and scripts has been installed or not. We don’t have to dig into much about how this part works.

The logical part is in the second part of the plugin code, which contains two parts. The first part is the core while second part simply attaches workon_cwd function from first part onto chpwd hook. So we only have to heed to the workon_cwd function.

The original code is shown below.

function workon_cwd { if [ ! $WORKON_CWD ]; then WORKON_CWD=1 # Check if this is a Git repo # Get absolute path, resolving symlinks PROJECT_ROOT=”${PWD:A}” while [[ “$PROJECT_ROOT” != “/” ! -e “$PROJECT_ROOT/.venv” ]]; do PROJECT_ROOT=”${PROJECT_ROOT:h}” done if [[ “$PROJECT_ROOT” == “/” ]]; then PROJECT_ROOT=”.” fi # Check for virtualenv name override if [[ -f “$PROJECT_ROOT/.venv” ]]; then ENV_NAME=`cat “$PROJECT_ROOT/.venv”` elif [[ -f “$PROJECT_ROOT/.venv/bin/activate” ]];then ENV_NAME=”$PROJECT_ROOT/.venv” elif [[ “$PROJECT_ROOT” != “.” ]]; then ENV_NAME=”${PROJECT_ROOT:t}” else ENV_NAME=”” fi if [[ “$ENV_NAME” != “” ]]; then # Activate the environment only if it is not already active if [[ “$VIRTUAL_ENV” != “$WORKON_HOME/$ENV_NAME” ]]; then if [[ -e “$WORKON_HOME/$ENV_NAME/bin/activate” ]]; then workon “$ENV_NAME” export CD_VIRTUAL_ENV=”$ENV_NAME” elif [[ -e “$ENV_NAME/bin/activate” ]]; then source $ENV_NAME/bin/activate export CD_VIRTUAL_ENV=”$ENV_NAME” fi fi elif [[ -n $CD_VIRTUAL_ENV -n $VIRTUAL_ENV ]]; then # We’ve just left the repo, deactivate the environment # Note: this only happens if the virtualenv was activated automatically deactivate unset CD_VIRTUAL_ENV fi unset PROJECT_ROOT unset WORKON_CWD fi }

From line 6 to 12, it checks the project root folder with the criteria is whether there is a .venv in folder. If not, the PROJECT_ROOT variable is . to represent current folder.

From line 13 to 22, it checks whether there is a .venv file to overwrite virtualenv name from PROJECTROOT folder to parent folder recursively. The issue here is that variable ENVNAME will be set to an empty string if .venv file could not be found. Then in the line 23, the ENV_NAME check will fail because of this, and the plugin fails to activate virtualenv.

The .venv is used to determine whether the folder has virtualenv instead of overwriting the virtualenv name. So we could change the code to set ENV_NAME to current folder name when .venv file is not found.

# Check for virtualenv name override if [[ -f “$PROJECT_ROOT/.venv” ]]; then ENV_NAME=`cat “$PROJECT_ROOT/.venv”` elif [[ -f “$PROJECT_ROOT/.venv/bin/activate” ]];then ENV_NAME=”$PROJECT_ROOT/.venv” elif [[ “$PROJECT_ROOT” != “.” ]]; then ENV_NAME=”${PROJECT_ROOT:t}” else FOLDER_NAME=`pwd` ENV_NAME=`basename “$FOLDER_NAME”` fi

Here we seems have set everything correct. But when we have sub-folders in virtualenv root folder, checking ENV_NAME would fail and then deactivate the environment because sub-folder name doesn’t match any virtualenv name. The new logical would illustrate as below.

The code with above logic is shown below.

# Activate the environment only if it is not already active if [[ “$VIRTUAL_ENV” != “$WORKON_HOME/$ENV_NAME” ]]; then if [[ -e “$WORKON_HOME/$ENV_NAME/bin/activate” ]]; then workon “$ENV_NAME” export CD_VIRTUAL_ENV=`pwd` elif [[ -e “$ENV_NAME/bin/activate” ]]; then source $ENV_NAME/bin/activate export CD_VIRTUAL_ENV=`pwd` elif [[ `pwd` != “$CD_VIRTUAL_ENV”* ]]; then # We’ve just left the repo, deactivate the environment # Note: this only happens if the virtualenv was activated automatically deactivate unset CD_VIRTUAL_ENV fi fi

Here we can say the new optimized virtualenvwrapper plugin for zsh has finished. It now can automatically activate & deactivate the environment when entering and leaving. The full file can be download from here.