zsh script to auto activate python virtual environments created by python2 virtualenv, virtualenvwrapper, and python3 venv

python-virtualenv.zsh

#
# this script checks for a .venv file or directory in your project path/root, and
# activates it automatically.. and deactivates it when you leave the path
#
# this may not work for you. it works for me.
#
# based on the virtualenvwrapper plugin from oh-my-zsh
#

function python_virtualenv_activate {

  # Check if this is a Git repo
  local GIT_REPO_ROOT=""
  local GIT_TOPLEVEL="$(git rev-parse --show-toplevel 2> /dev/null)"
  if [[ $? == 0 ]]; then
    GIT_REPO_ROOT="$GIT_TOPLEVEL"
  fi

  # Get absolute path, resolving symlinks
  local PROJECT_ROOT="${PWD:A}"
  while [[ "$PROJECT_ROOT" != "/" && ! -e "$PROJECT_ROOT/.venv" \
           && ! -d "$PROJECT_ROOT/.git"  && "$PROJECT_ROOT" != "$GIT_REPO_ROOT" ]]; do
    PROJECT_ROOT="${PROJECT_ROOT:h}"
  done
  if [[ "$PROJECT_ROOT" == "/" ]]; then
    PROJECT_ROOT="."
  fi

  local ENV_ROOT=""

  # look for a file or directory named .venv in the PROJECT_ROOT directory.
  # if it's a file, it contains a virtual environment name, and we should
  # look in WORKON_HOME plus other common virtualenv home directories for
  # an activate script ... if it's a directory, the virtualenv is in the
  # PROJECT_ROOT and we should check in there for an activate script...
  if [[ -f "$PROJECT_ROOT/.venv/bin/activate" ]];then
    ENV_NAME="${PROJECT_ROOT:t}"
    ENV_ROOT="$PROJECT_ROOT/.venv/"
  elif [[ -f "$PROJECT_ROOT/.venv" ]]; then
    ENV_NAME="$(cat "$PROJECT_ROOT/.venv")"
    if [[ -f "$WORKON_HOME/$ENV_NAME/bin/activate" ]]; then
      ENV_ROOT="$WORKON_HOME/$ENV_NAME"
    elif [[ -f "$HOME/.python/$ENV_NAME/bin/activate" ]]; then
      ENV_ROOT="$HOME/.python/$ENV_NAME"
    elif [[ -f "$HOME/.virtualenv/$ENV_NAME/bin/activate" ]]; then
      ENV_ROOT="$HOME/.virtualenv/$ENV_NAME"
    else
      ENV_ROOT=""
      ENV_NAME=""
    fi
  else
    ENV_NAME=""
    ENV_ROOT=""
  fi

  if [[ "$ENV_NAME" != "" ]]; then
    # Activate the environment only if it is not already active
    if [[ "$VIRTUAL_ENV" != "$ENV_ROOT" ]]; then
      if [[ -e "$ENV_ROOT/bin/activate" ]]; then
        source $ENV_ROOT/bin/activate && export CD_VIRTUAL_ENV="$ENV_ROOT"
      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
}

# Append workon_cwd to the chpwd_functions array, so it will be called on cd
# http://zsh.sourceforge.net/Doc/Release/Functions.html
if ! (( $chpwd_functions[(I)python_virtualenv_activate] )); then
  chpwd_functions+=(python_virtualenv_activate)
fi