scripts/misc/venv

#!/bin/zsh
#dep:fd fzf

source ${$(realpath "$0"):h:h}/lib/common.zsh

zparseopts -D -E \
	c=cdvenv -cdvenv=cdvenv \
	h=help -help=help \
	n=newvenv -new=newvenv \
	t=tmsu -tmsu=tmsu \
	p=pyenv -pyenv=pyenv \
	f=find -find=find \
	i=ignore -ignore-local=ignore \
	s=save -save=save

if [[ $help ]]; then
	<<-HELP
	Usage: ${0:t} [-c] [DIR]
	       ${0:t} -t

	If DIR contains bin/activate, start a shell with that virtualenv.
	Otherwise recursively look up virtualenvs in DIR (by looking for dirs ending
	in "venv") and select one with fzf.

	OPTIONS:
	 -c, --cdvenv  Change to the venv directory instead of going back to the 
				   current working directory (not for pyenv virtualenvs)
	 -t, --tmsu    list only venvs named "venv" in a directory with tmsu tag lang=python
	 -p, --pyenv   Use pyenv virtualenvs
	 -f, --find    Search for virtualenvs by looking for directories ending with 'venv'

	    If neither -p or -f are given, the default behaviour is to use both

	 -s, --save    Store selected venv in .venv in current directory. Will be used 
	               automatically, when called in this dir again.
	 -i, --ignore-local
	               Ignore .venv files.
	HELP
	exit
fi

ORIGDIR=$PWD

venvsh() {
	cd $1
	. bin/activate
	if [[ -z $cdvenv ]]; then
		cd $ORIGDIR
	fi
	exec $SHELL
}

pyenvsh() {
	export PYENV_VIRTUALENV_DISABLE_PROMPT=1
	eval "$(pyenv init -)"
	eval "$(pyenv virtualenv-init -)"
	pyenv activate "$1"
	exec $SHELL
}

envsh() {
	local venv=$1
	if [[ $venv =~ pyenv\!* ]]; then
		pyenvsh ${venv##pyenv!}
	else
		venvsh $venv
	fi
}

if [[ -z $ignore && -e .venv ]]; then
	venv=$(<.venv)
	echo "Local .venv found (-i to ignore)."
	echo "Using: $venv"
	envsh $venv
fi


if [[ -n $tmsu ]]; then
	venv=$(
		for i in $(tmsu files --directory lang=python); do
			[[ -d $i/venv ]] && echo $i/venv
		done | fzf)
	if [[ -n $venv ]]; then
		exec venvsh $venv
	else
		exit 1
	fi
fi

if [[ $newvenv ]]; then
	python -mvenv ${1:-venv}
	exec venvsh ${1:-venv}
fi

if [[ -n $1 && -e $1/bin/activate ]]; then
	exec venvsh $1
fi

pyenvs() {
	if exists pyenv && pyenv commands | grep -q 'virtualenvs'; then
		pyenv virtualenvs --bare --skip-aliases | sed -e 's/^/pyenv!/'
	fi
}

if [[ $pyenv && ! $find ]]; then
	venv=$(pyenvs | fzf)
elif [[ ! $pyenv && $find ]]; then
	venv=$(fd -I -t d 'venv$' $1 | fzf)
else
	venv=$(cat <(fd -I -t d 'venv$' $1) <(pyenvs) | fzf)
fi

if [[ -n $venv ]]; then
	if [[ -n $save ]]; then
		echo $venv > .venv
	fi
	envsh $venv
fi