Add shell functions for Python venvs

It is an ugly system, and it seems it is practically inescapable. Truly
an homage to the quality of the programming language and the devs who
love it without question.

3 files changed, 61 insertions, 1 deletions

diff --git a/.bashrc b/.bashrc
index f526efa..ca3a9d3 100644
--- a/.bashrc
+++ b/.bashrc
@@ -72,6 +72,9 @@ __precmd()
 	PS1+="\001$SHELLC_F_LIGHTGRAY\002:\001$SHELLC_RESET_COLOR_F\002"
 	PS1+="\001$SHELLC_F_BLUE\002\w\001$SHELLC_RESET_COLOR_F\002"
 
+	# Add venv indicator
+	PS1+="\001$SHELLC_F_CYAN\002#$(awk -F/ '{ print $(NF-1) }' <<< "$PYTHON_VENV_DIR")\001$SHELLC_RESET_COLOR_F\002"
+
 	# Add git status
 	if git rev-parse --is-inside-work-tree &> /dev/null
 	then
diff --git a/.config/shell/functions.d/venv b/.config/shell/functions.d/venv
new file mode 100644
index 0000000..50a91e9
--- /dev/null
+++ b/.config/shell/functions.d/venv
@@ -0,0 +1,48 @@
+#!/bin/sh
+
+venv() {
+	if ! command -v "venv_$1" | grep -q " function "
+	then
+		cat <<EOF
+Usage:
+	venv on
+	venv off
+
+Manage the python venv for a given project or directory.
+
+Commands:
+	on	Enable the venv for \$PWD
+	off	Deactivate the current venv
+EOF
+
+		return 1
+	fi
+
+	"venv_$1"
+}
+
+venv_on() {
+	export PYTHON_VENV_DIR="$PWD/.venv"
+	export VIRTUAL_ENV_DISABLE_PROMPT=1
+
+	if [ ! -d "$PYTHON_VENV_DIR" ]
+	then
+		printf "Initializing venv directory at %s\n" "$PYTHON_VENV_DIR"
+		python3 -m venv "$PYTHON_VENV_DIR"
+	fi
+
+	. "$PYTHON_VENV_DIR/bin/activate"
+}
+
+venv_off() {
+	if [ -z "$PYTHON_VENV_DIR" ]
+	then
+		printf "No venv active?\n"
+		return 1
+	fi
+
+	deactivate
+
+	unset PYTHON_VENV_DIR
+	unset VIRTUAL_ENV_DISABLE_PROMPT
+}
diff --git a/.zshrc b/.zshrc
index 63a89a1..347f3c0 100644
--- a/.zshrc
+++ b/.zshrc
@@ -126,6 +126,14 @@ function precmd() # {{{
 		sshd|*/sshd|mosh*) HOSTNAME_COLOR="%F{6}" ;;
 	esac
 	# }}}
+
+	# Color last directory name in case of python venv
+	PROMPT_PATH='%F{4}%~%k'
+
+	if [[ -n "$PYTHON_VENV_DIR" ]]
+	then
+		PROMPT_PATH+='%F{6}#'"$(awk -F/ '{ print $(NF-1) }' <<< "$PYTHON_VENV_DIR")"
+	fi
 }
 # }}}
 
@@ -162,7 +170,8 @@ PROMPT='%(!.%F{1}.%F{2})%n%F{8}'
 PROMPT=$PROMPT'@'
 PROMPT=$PROMPT'${HOSTNAME_COLOR}%M%F{8}'
 PROMPT=$PROMPT':'
-PROMPT=$PROMPT'%F{4}%~%k${PGIT}'
+PROMPT=$PROMPT'${PROMPT_PATH}'
+PROMPT=$PROMPT'%k${PGIT}'
 PROMPT=$PROMPT' '
 PROMPT=$PROMPT'%(?.%F{7}.%F{1})${EXITCODE}%k'
 PROMPT=$PROMPT' ${vi_mode}»%b %f${CURSOR_STYLE}'