dotfiles/functions/python.zsh

mkvenv(){
    local source_venv force_create venv_name
    while getopts "sf" opt; do
        case $opt in
            s|source) source_venv=y ;;
            f|force) force_create=y ;;
            *) echo "Unknown option: $opt"; return 1 ;;
        esac
    done

    shift $(($OPTIND-1))
    if [[ -n $1 ]]; then
        venv_name=$1
        shift
    else
        venv_name=.venv
    fi
    # any other arguments are passed on to the `venv` module as flags/arguments


    if [[ -e ${venv_name} ]]; then
        if [[ -n ${force_create} ]]; then
            if [[ -n $VIRTUAL_ENV ]]; then # deactivate venv if enabled
                deactivate && warning "Deactivated existing virtualenv" || red "Could not run deactivate";
            fi
            rm -rf ${venv_name} && warning "Deleted existing virtualenv."
        else
            error "$(color $BOLD .venv) already exists. Run with -f to recreate it." >&2
            return 1
        fi
    fi

    info "Creating virtualenv ${venv_name}... "

    if ! python -m virtualenv $@ ${venv_name} ; then
        error "Failed to create virtualenv. Is virtualenv installed? Try:\n $ pip install virtualenv"
        return 1
    fi

    if [[ -z ${source_venv} ]]; then
        echo -n " done. Enable? [Y/n] "
        read source_venv
    else
        echo ""
    fi

    if [[ $source_venv != "n" ]]; then
        source ${venv_name}/bin/activate && \
        echo -e "  → $(color $BOLD Enabled!) 🐍 $(color $BOLD$PURPLE $(python --version | cut -d " " -f2 )) ($(color $BOLD$GREEN $(pip --version | cut -d " " -f -2)))"
    fi
}

function enable_venv() {
    local root
    local activate

    root=$(git rev-parse --show-toplevel 2>/dev/null)
    if [[ -z $root ]]; then root=$PWD; fi
    activate="$root/.venv/bin/activate"
    if [[ -f "$activate" ]]; then
        source "$activate"
    fi
}

function disable_venv() {
    if [[ -z ${VIRTUAL_ENV} ]]; then
        return
    fi

    if [[ ! "$PWD" =~ $(dirname $VIRTUAL_ENV)}* ]]; then
        echo "Deactivating venv..."
        deactivate
    fi
}


if [[ -n $DOTFILES_AUTOSWITCH_VIRTUALENV ]]; then
    add-zsh-hook chpwd enable_venv
    add-zsh-hook chpwd disable_venv

    enable_venv
fi