.pyvenvwrapper

#!/bin/bash

################################################################################
#   Python Venv Wrapper
################################################################################

if [[ -z "$VENV_HOME" ]]; then
    export VENV_HOME="$HOME/.venv" # default
fi

if [[ -z "$VENV_PROJECT_FOLDER" ]]; then
    export VENV_PROJECT_FOLDER="$HOME/Projects" # default
fi

echoerr() {
    printf "E: %s\n" "$*" >&2;
}

[[ -d "$VENV_HOME" ]] || mkdir "$VENV_HOME"

if [[ ! -d "$VENV_PROJECT_FOLDER" ]]; then
    echo -n "$VENV_PROJECT_FOLDER does not exist. Create? [y/n]: "
    read -r -n 1 ans
    echo
    if [[ $ans == "y" ]] || [[ $ans == "Y" ]]; then
        if mkdir "$VENV_PROJECT_FOLDER"; then
            echo "Created Project folder at $VENV_PROJECT_FOLDER"
        else
            echoerr "Failed creating $VENV_PROJECT_FOLDER"
        fi
    fi
fi

if [[ ! -d "$VENV_PROJECT_FOLDER" ]] && [[ -z "$VENV_NO_PROJ" ]]; then
    echoerr "No project folder available!"
    echoerr "    Please create a folder at $VENV_PROJECT_FOLDER"
    echoerr "    or export a different folder name to variable VENV_PROJECT_FOLDER"
    echoerr "    or set a variable VENV_NO_PROJ=1 in order to skip this check entirely. (workon and mkproj may not work)."
    return 1
fi

lsvenv() {
    command ls "$VENV_HOME" | tr '\n' '\0' | xargs -0 -n 1 basename 2>/dev/null
}

activate() {
    if [[ $# -le 0 ]]; then
        echoerr "No arguments provided."
        return 1
    fi

    if [[ ! -d "$VENV_HOME/$1" ]]; then
        echoerr "Environment '$VENV_HOME/$1' does not contain an activate script."
        return 1
    fi

    [[ -n "$VIRTUAL_ENV" ]] && deactivate

    # shellcheck source=/dev/null
    source "${VENV_HOME:?}/$1/bin/activate"
}

workon() {
    if [[ ! -d "$VENV_PROJECT_FOLDER" ]]; then
        echoerr "The root projects folder $VENV_PROJECT_FOLDER does not exist."
        return 1
    fi

    if [[ $# -le 0 ]]; then
        echoerr "No arguments provided."
        return 1
    fi

    if [[ ! -d "$VENV_PROJECT_FOLDER/$1" ]]; then
        echoerr "The project $VENV_PROJECT_FOLDER/$1 does not exist."
        return 1
    fi

    activate "$@" && cd "${VENV_PROJECT_FOLDER:?}/$1" || exit
}

mkproj() {
    if [[ ! -d "$VENV_PROJECT_FOLDER" ]]; then
        echoerr "The root projects folder $VENV_PROJECT_FOLDER does not exist."
        return 1
    fi

    if [[ $# -le 0 ]]; then
        echoerr "No arguments provided."
        return 1
    fi

    if [[ -d "$VENV_PROJECT_FOLDER/$1" ]]; then
        echoerr "The project $VENV_PROJECT_FOLDER/$1 already exists."
        return 1
    fi

    if mkvenv "$@"; then
        mkdir "$VENV_PROJECT_FOLDER/$1"
        echo "Project folder created at $VENV_PROJECT_FOLDER/$1."
        cd "${VENV_PROJECT_FOLDER:?}/$1" || exit
    fi
}

mkvenv() {
    if [[ $# -le 0 ]]; then
        echoerr "No arguments provided."
        return 1
    fi

    VENV_LOCATION=${VENV_HOME:?}/$1

    if [[ -d "$VENV_LOCATION" ]]; then
        echoerr "$1 already exists in $VENV_HOME."
        return 1
    fi

    [[ -n "$VIRTUAL_ENV" ]] && deactivate
    python3 -m venv "${VENV_LOCATION:?}"

    # shellcheck source=/dev/null
    if source "${VENV_LOCATION:?}/bin/activate"; then
        echo "Python venv created at $VENV_LOCATION.";
    fi

    [[ -z $VENV_SKIP_PIP_UPGRADE ]] && pip install --upgrade pip
}

rmvenv() {
    if [[ $# -le 0 ]]; then
        echoerr "No arguments provided."
        return 1
    fi

    VENV_LOCATION=${VENV_HOME:?}/$1

    if [[ ! -e "$VENV_HOME/$1" ]]; then
        echoerr "Environment '$VENV_LOCATION' does not exist."
        return 1
    fi

    [[ "$VIRTUAL_ENV" == "$VENV_LOCATION" ]] && deactivate
    if rm -r "$VENV_LOCATION" ; then
        echo "Python venv removed at $VENV_LOCATION."
    fi
}