Support zsh and deactivate when leaving environment directories

README.md

@@ -8,7 +8,7 @@ This functionality was inspired by [conda auto activate](https://github.com/sott
 
 ## Install
 
-To install add this line to your .bashrc or .bash-profile:
+To install add this line to your .bashrc, .bash-profile, or .zshrc:
 
        source /path/to/conda_auto_env.sh
 

conda_auto_env.sh

@@ -15,20 +15,42 @@ function conda_auto_env() {
   if [ -e "environment.yml" ]; then
     # echo "environment.yml file found"
     ENV=$(head -n 1 environment.yml | cut -f2 -d ' ')
-    # Check if you are already in the environment
-    if [[ $PATH != *$ENV* ]]; then
-      # Check if the environment exists
-      source activate $ENV
-      if [ $? -eq 0 ]; then
-        :
-      else
-        # Create the environment and activate
-        echo "Conda env '$ENV' doesn't exist."
+    # echo "Check if the environment is already active."
+    if [[ $PATH != */envs/$ENV/* ]]; then
+      # echo "Attempt to activate environment."
+      CONDA_ENVIRONMENT_ROOT="" #For spawned shells
+      conda activate $ENV
+      # echo "Set root directory of active environment."
+      CONDA_ENVIRONMENT_ROOT="$(pwd)"
+      if [ $? -ne 0 ]; then
+        # Create the environment and activate.
+        echo "Conda environment '$ENV' doesn't exist: Creating."
         conda env create -q
-        source activate $ENV
+        conda activate $ENV
       fi
     fi
+  elif [[ $PATH == */envs/* ]]\
+    && ([[ $(pwd) != $CONDA_ENVIRONMENT_ROOT ]]\
+      || [[ $(pwd) != $CONDA_ENVIRONMENT_ROOT/* ]])
+  then
+    # echo "Deactivate active environment since we are no longer among its subdirectories."
+    CONDA_ENVIRONMENT_ROOT=""
+    # echo "Before conda deactivate"
+    conda deactivate
+    # echo "After conda deactivate"
   fi
 }
 
-export PROMPT_COMMAND=conda_auto_env
+# Check active shell.
+if [[ $(ps -p$$ -ocommand=) == "-zsh" ]]; then
+  # For zsh, use the chpwd hook.
+  autoload -U add-zsh-hook
+  add-zsh-hook chpwd conda_auto_env
+  # Run for present directory as it does not fire the above hook.
+  conda_auto_env
+  # More aggressive option in case the above hook misses some use case:
+  #precmd() { conda_auto_env; }
+else
+  # For bash, no hooks and we rely on the env. var. PROMPT_COMMAND:
+  export PROMPT_COMMAND=conda_auto_env
+fi

conda_auto_env_remote.sh

@@ -18,39 +18,61 @@ function conda_auto_env_remote() {
   if [ -e "environment.yml" ]; then
     # echo "environment.yml file found"
     ENV=$(head -n 1 environment.yml | cut -f2 -d ' ')
-    # Check if you are already in the environment
-    if [[ $PATH != *$ENV* ]]; then
-      # Check if the environment exists
-      source activate $ENV
-      if [ $? -eq 0 ]; then
-        :
-      else
-        # Create the environment and activate
-        echo "Conda env '$ENV' doesn't exist."
+    # Check if the environment is already active.
+    if [[ $PATH != */envs/$ENV/* ]]; then
+      # Attempt to activate environment.
+      CONDA_ENVIRONMENT_ROOT="" #For spawned shells
+      conda activate $ENV
+      # Set root directory of active environment.
+      CONDA_ENVIRONMENT_ROOT="$(pwd)"
+      if [ $? -ne 0 ]; then
+        # Create the environment and activate.
+        echo "Conda environment '$ENV' doesn't exist: Creating."
         conda env create -q
-        source activate $ENV
+        conda activate $ENV
       fi
     fi
   fi
   if [ -e "environment-remote.yml" ]; then
     # echo "environment.yml file found"
     ENV=$(sed -n '1p' environment-remote.yml | cut -f2 -d ' ')
     CHANNEL=$(sed -n '2p' environment-remote.yml | cut -f2 -d ' ')
-    # Check if you are already in the environment
-    if [[ $PATH != *$ENV* ]]; then
-      # Check if the environment exists
-      source activate $ENV
-      if [ $? -eq 0 ]; then
-        :
-      else
-        # Create the environment and activate
+    # Check if the environment is already active.
+    if [[ $PATH != */envs/$ENV/* ]]; then
+      # Attempt to activate environment.
+      CONDA_ENVIRONMENT_ROOT="" #For spawned shells
+      conda activate $ENV
+      # Set root directory of active environment.
+      CONDA_ENVIRONMENT_ROOT="$(pwd)"
+      if [ $? -ne 0 ]; then
+        # Create the environment and activate.
         echo "Conda env '$ENV' doesn't exist."
         REMOTE=$CHANNEL'/'$ENV
         conda env create $REMOTE -q
-        source activate $ENV
+        conda activate $ENV
       fi
     fi
   fi
+  # Deactivate active environment if we are no longer among its subdirectories.
+  if [[ $PATH = */envs/* ]]\
+    && ([[ $(pwd) != $CONDA_ENVIRONMENT_ROOT ]]\
+      || [[ $(pwd) != $CONDA_ENVIRONMENT_ROOT/* ]])
+  then
+    CONDA_ENVIRONMENT_ROOT=""
+    conda deactivate
+  fi
 }
 
-export PROMPT_COMMAND=conda_auto_env_remote
+# Check active shell.
+if [[ $(ps -p$$ -ocommand=) == "-zsh" ]]; then
+  # For zsh, use the chpwd hook.
+  autoload -U add-zsh-hook
+  add-zsh-hook chpwd conda_auto_env
+  # Run for present directory as it does not fire the above hook.
+  conda_auto_env
+  # More aggressive option in case the above hook misses some use case:
+  #precmd() { conda_auto_env; }
+else
+  # For bash, no hooks and we rely on the env. var. PROMPT_COMMAND:
+  export PROMPT_COMMAND=conda_auto_env
+fi