Setting env vars automatically in Virtual Environments

As part of my development flow on Airflow I use virtualenvwrapper and a long time ago I set up the ability to automatically set environment variables whenever the virtualenv is activated, and to unset them when leaving. (I think this trick would probably work with a few small changes for other venv tools and isn’t that specific to virtualenvwrapper).

This trick lets me create an “env” file so that I can easily switch between different configs as I change virtualenvs, super useful as I switch back and forth between versions or have multiple envs in parallel when I’m testing performance etc.

For example my “main” dev virtualenv has ~.virtualenvs/airflow/share/auto-env.zsh containing:

AIRFLOW__CORE__SQL_ALCHEMY_CONN=postgresql:///airflow
AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=postgresql:///airflow

How it works

In ~/.virtualenvs (the path I have configured to keep my environments) there are a set of “hook scripts”, and I tweaked two of them in particular:

postactivate:

This one is fairly straight forward, but makes use of set -a which automatically exports any variable set until set +a is encountered. This just means the env file doesn’t need export on every line — just a bit simpler. (set -a works on Bash as well as ZSH).

#!/usr/bin/zsh
# This hook is sourced after every virtualenv is activated.

if [[ -e "$VIRTUAL_ENV/share/auto-env.zsh" ]]; then
set -a
source "$VIRTUAL_ENV/share/auto-env.zsh"
set +a
fi

predeactivate:

This script is a bit more complex, but it goes and “parses” the env file (in a very simplistic manner) looking for the name of all the env vars we set, and then unsets them.

#!/usr/bin/zsh
# This hook is sourced before every virtualenv is deactivated.

autoenv="$VIRTUAL_ENV/share/auto-env.zsh"

if [[ -e "$VIRTUAL_ENV/share/auto-env.zsh" ]]; then
for variable in $(grep -oP '^[a-zA-Z_][a-zA-Z-1-9_]*(?==)' "$autoenv"); do
    unset "${variable}"
done
fi