Script to run Pip across multiple Virtual Environments

pip_virtualenvs.sh

#!/bin/bash
######################################################################
# Useful script when you want to run `pip` in virtual environments   #
# you have in a common directory with a provided Python package.     #
# Note: Requires VirtualWrapper to be managing your environments     #
#                                                                    #
# Usage:                                                             #
# ./pip_virtualenvs.sh [-i | -u | -z] pkg_name [-a | -l] [env1 env2] #
#                                                                    #
# Example:                                                           #
# ./pip_virtualenvs.sh requests -l my_cool_env1 my_cool_env2         #
#                                                                    #
# Options:                                                           #
# Pip Action Flags                                                   #
# -i  OR  --install:    pip install pkg (default action)             #
#                                                                    #
# -u  OR  --upgrade:    pip install pkg --upgrade                    #
#                                                                    #
# -z  OR  --uninstall:  pip uninstall pkg                            #
#                                                                    #
# Virtual Environment Flags                                          #
# -a  OR  --all:        runs pip in every virtualenv on the system   #
#                                                                    #
# -l  OR  --list:       list the envs to pip in (default action)     #
######################################################################
function setup () {
    # save current directory for later
    STARTING_DIR=$(pwd);

    # default pip action is install
    PIP_ACTION='install'

    # source virtualwrapper.sh, just in case
    source $(which virtualenvwrapper.sh)

    # default virtualenvwrapper directory is "$HOME/.virtualenvs"
    if [ -d "$VIRTUALENVWRAPPER_HOOK_DIR" ]; then
        VIRTUALENVS_DIR="$VIRTUALENVWRAPPER_HOOK_DIR"
    else
        _handle_error "envdir_err"
    fi
}

function handle_options () {
    # check for pip action option flags, if any
    # NOTE: we shift if there are flags so $1
    # is always the right parameter to check
    # NOTE: only ONE of the below flags should be passed
    case "$1" in
        -i|--install) PIP_ACTION='install'; shift ;;

        # note PIP_ADDITION here, since upgrade syntax is
        # `pip install $pkg --upgrade`, not `pip upgrade $pkg`
        -u|--upgrade) PIP_ADDITION='--upgrade'; shift ;;

        -z|--uninstall) PIP_ACTION='uninstall'; shift ;;

        -*) _handle_error "action_err" ;;

    esac

    # package we want to use with pip
    PIP_PACKAGE=$1; shift;

    # NOTE: we shift if there are flags so $1
    # is always the right parameter to check
    if [ -z "$1" ]; then
        _handle_error "no_envs_err"
    fi

    VIRTUALENVS_ARRAY=()
    while [ ! -z "$1" ]; do
        case "$1" in
            # run pip in ALL virtualenvs on the system
            -a|--all) _mk_all_virtualenv_array; return ;;

            # list specific virtualenvs to pip in
            # take second arg, since $1 will be -l or --list
            # then shift twice when done to move onto the next arg
            -l|--list) _mk_selected_virtualenv_array $2; shift; shift ;;

            -*) _handle_error "action_err" ;;

            # defaults to list selection if no flag provided
            *) _mk_selected_virtualenv_array $1; shift ;;
        esac
    done
}

function _mk_all_virtualenv_array () {
    # change into the virtualenvs directory
    cd $VIRTUALENVS_DIR

    # set nullglob to avoid including "*/" if no subdirectories available
    shopt -s nullglob

    # enter the results into an array
    VIRTUALENVS_ARRAY=(*/)

    # turn off nullglob to avoid messing with stuff later
    shopt -u nullglob
}

function _mk_selected_virtualenv_array () {
    ENV="$1"

    if [ ! -d "$VIRTUALENVS_DIR/$ENV/" ]; then
        _handle_error "virtualenv_err" "$1"
    fi;

    # add ENV to VIRTUALENVS_ARRAY
    VIRTUALENVS_ARRAY+=("$ENV")

}

function main () {
    # validate function arguments
    _validate_func

    # update each virtual environment
    for ENV_DIR in ${VIRTUALENVS_ARRAY[@]}; do
        _update_virtualenv "$ENV_DIR";
    done

    # return to the original working directory
    cd "$STARTING_DIR"

    exit 0
}

# set defaults for key variables if none are provided
function _validate_func () {
    # error handler if no pip package is provided
    if [ -z "$PIP_PACKAGE" ]; then
        _handle_error "pkg_err"
    fi

    if [ -z "$VIRTUALENVS_ARRAY" ]; then
        _handle_error "no_envs_err"
    fi

}

function _update_virtualenv () {
    ENV_DIR=$1
    ACTIVATION_FILE="$VIRTUALENVS_DIR/$ENV_DIR/bin/activate"
    echo $ACTIVATION_FILE

    if [ ! -e "$ACTIVATION_FILE" ]; then
        echo "Missing activation file: $ENV_DIR. Moving on..."
        return
    fi

    # activate the relevant virtualenv
    echo "Activating virtualenv for directory: $ENV_DIR..."
    source "$ACTIVATION_FILE"

    # run the specified pip action with the pip package
    pip $PIP_ACTION $PIP_PACKAGE $PIP_ADDITION

    # deactivate the virtual environment
    deactivate

    #print success message
    echo "###############################################################"
    echo "$PIP_ACTION complete in: $ENV_DIR! Deactivating virtualenv..."
    echo "###############################################################"
    echo
}

function _handle_error () {
    # handle error based on error code provided
    ERR_CODE=$1
    case "$ERR_CODE" in
        "pkg_err")
            ERR_MSG="Must provide pip package!"
        ;;

        "envdir_err")
            ERR_MSG="Virtualenv directory missing! \
            Are you using virtualenvwrapper?"
        ;;

        "action_err")
            ERR_MSG="Pip action not recognized!"
        ;;

        "no_envs_err")
            ERR_MSG="Must provide virtualenvs to pip in!"
        ;;

        "virtualenv_err")
            ERR_MSG="$2 is not a valid virtualenv, aborting! Available envs: \
            `echo; ls $VIRTUALENVS_DIR`"
        ;;

        *)
            ERR_MSG="Unknown error!"
        ;;
    esac


    # print error message
    echo "Error: $ERR_MSG"
    echo
    echo 'Usage: ./pip_virtualenvs.sh [-i | -u | -z] pkg_name [-a | -l] [env1 env2 env3]...'
    echo
    echo 'Options:'
    echo 'Pip Action Flags'
    echo '-i  OR  --install:    pip install pkg (default action)'
    echo
    echo '-u  OR  --upgrade:    pip install pkg --upgrade'
    echo
    echo '-z  OR  --uninstall:  pip uninstall pkg'
    echo
    echo 'Virtual Environment Flags'
    echo '-a  OR  --all:        runs pip in every virtualenv on the system'
    echo
    echo '-l  OR  --list:       list the virtualenvs to run pip in (default action)'
    echo '                      example: ./pip_virtualenvs.sh pkg -l my_cool_env1 my_cool_env2'
    echo


    # deactivate any active virtualenv
    if [ ! -z "$VIRTUAL_ENV" ]; then
        deactivate
    fi

    # exit with an error code
    exit 1
}

##################
# MAIN PROCEDURE #
##################
setup
handle_options $*
main