Better process for zsh virtualenvwrapper plugin

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.