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}'