From 08e39a587fa130f04077201d050e72b2130d637e Mon Sep 17 00:00:00 2001
From: Matthew Crenshaw <sgt@squig.gs>
Date: Tue, 30 May 2017 20:26:18 -0400
Subject: [PATCH] Python virtualenv auto workon cwd hook (#1300)

* Command to activate virtualenvs on changing directory

Enables automatic activation of a virtualenv when jumping into a directory
This is done by looking in the current directory for a file or directory
named `.venv`. If it's not found in the current directory its parents will
also be examined. `.venv` can be either:
    1. A file containing the name of a virtualenv found in $WORKON_HOME
    2. A directory containing bin/activate (meaning that the directory is
       assumed to be a virtualenv.)
If $WORKON_HOME is set it is assumed that virtualenvwrapper is installed
and the `workon` command will be issued.
---
 modules/python/init.zsh | 48 +++++++++++++++++++++++++++++++++++++++++
 runcoms/zpreztorc       |  7 ++++++
 2 files changed, 55 insertions(+)

diff --git a/modules/python/init.zsh b/modules/python/init.zsh
index 4f67a081..47f7b5a1 100644
--- a/modules/python/init.zsh
+++ b/modules/python/init.zsh
@@ -34,6 +34,48 @@ if (( ! $+commands[python] && ! $+commands[pyenv] )); then
   return 1
 fi
 
+function _python-workon-cwd {
+  # 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
+  # Check for virtualenv name override
+  local ENV_NAME=""
+  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}"
+  fi
+  if [[ -n $CD_VIRTUAL_ENV && "$ENV_NAME" != "$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
+  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
+  fi
+}
+
 # Load virtualenvwrapper into the shell session, unless requested not to
 if zstyle -T ':prezto:module:python' skip-virtualenvwrapper-init; then
   # Set the directory where virtual environments are stored.
@@ -65,6 +107,12 @@ if (( $+commands[pip] )); then
   source "$cache_file"
 
   unset cache_file
+
+  # Load auto workon cwd hook
+  if zstyle -t ':prezto:module:python:virtualenv' auto-switch 'yes'; then
+    # Auto workon when changing directory
+    add-zsh-hook chpwd _python-workon-cwd
+  fi
 fi
 
 #
diff --git a/runcoms/zpreztorc b/runcoms/zpreztorc
index 2ac6c58c..690add7a 100644
--- a/runcoms/zpreztorc
+++ b/runcoms/zpreztorc
@@ -107,6 +107,13 @@ zstyle ':prezto:module:prompt' theme 'sorin'
 # Auto switch the Ruby version on directory change.
 # zstyle ':prezto:module:ruby:chruby' auto-switch 'yes'
 
+#
+# Python
+#
+
+# Auto switch the Python virtualenv on directory change.
+# zstyle ':prezto:module:python:virtualenv' auto-switch 'yes'
+
 #
 # Screen
 #