zsh-autoswitch-virtualenv.plugin.zsh
15124 bytes
1export AUTOSWITCH_VERSION="3.9.0"
2export AUTOSWITCH_FILE=".venv"
3
4AUTOSWITCH_RED="\e[31m"
5AUTOSWITCH_GREEN="\e[32m"
6AUTOSWITCH_PURPLE="\e[35m"
7AUTOSWITCH_BOLD="\e[1m"
8AUTOSWITCH_NORMAL="\e[0m"
9
10VIRTUAL_ENV_DIR="${AUTOSWITCH_VIRTUAL_ENV_DIR:-$HOME/.virtualenvs}"
11
12function _validated_source() {
13 local target_path="$1"
14
15 if [[ "$target_path" == *'..'* ]]; then
16 (>&2 printf "AUTOSWITCH WARNING: ")
17 (>&2 printf "target virtualenv contains invalid characters\n")
18 (>&2 printf "virtualenv activation cancelled\n")
19 return
20 else
21 source "$target_path"
22 fi
23}
24
25
26function _virtual_env_dir() {
27 local venv_name="$1"
28 printf "%s/%s" "$VIRTUAL_ENV_DIR" "$venv_name"
29}
30
31
32function _python_version() {
33 local PYTHON_BIN="$1"
34 if [[ -f "$PYTHON_BIN" ]] then
35 # For some reason python --version writes to stderr
36 printf "%s" "$($PYTHON_BIN --version 2>&1)"
37 else
38 printf "unknown"
39 fi
40}
41
42
43function _autoswitch_message() {
44 if [ -z "$AUTOSWITCH_SILENT" ]; then
45 (>&2 printf "$@")
46 fi
47}
48
49
50function _is_valid_virtualenv() {
51 [[ -d "$1" ]] && [[ -f "$1/bin/activate" ]]
52}
53
54
55function _get_venv_type() {
56 local venv_dir="$1"
57 local venv_type="${2:-virtualenv}"
58 if [[ -f "$venv_dir/Pipfile" ]]; then
59 venv_type="pipenv"
60 elif [[ -f "$venv_dir/poetry.lock" ]]; then
61 venv_type="poetry"
62 elif [[ -f "$venv_dir/uv.lock" ]]; then
63 venv_type="uv"
64 elif [[ -f "$venv_dir/requirements.txt" || -f "$venv_dir/setup.py" ]]; then
65 venv_type="virtualenv"
66 fi
67 printf "%s" "$venv_type"
68}
69
70
71function _get_venv_name() {
72 local venv_dir="$1"
73 local venv_type="$2"
74 local venv_name="$(basename "$venv_dir")"
75
76 # clear pipenv from the extra identifiers at the end
77 if [[ "$venv_type" == "pipenv" ]]; then
78 venv_name="${venv_name%-*}"
79 fi
80
81 printf "%s" "$venv_name"
82}
83
84
85function _maybeworkon() {
86 local venv_dir="$1"
87 local venv_type="$2"
88 local venv_name="$(_get_venv_name $venv_dir $venv_type)"
89
90 local DEFAULT_MESSAGE_FORMAT="Switching to ${AUTOSWITCH_BOLD}%venv_type${AUTOSWITCH_NORMAL} project: ${AUTOSWITCH_BOLD}${AUTOSWITCH_PURPLE}%venv_name${AUTOSWITCH_NORMAL} ${AUTOSWITCH_GREEN}[%py_version]${AUTOSWITCH_NORMAL}"
91
92 # Don't reactivate an already activated virtual environment
93 if [[ -z "$VIRTUAL_ENV" || "$venv_dir" != "$VIRTUAL_ENV" ]]; then
94
95 if [[ ! -d "$venv_dir" ]]; then
96 printf "Unable to find ${AUTOSWITCH_PURPLE}$venv_name${AUTOSWITCH_NORMAL} virtualenv\n"
97 printf "If the issue persists run ${AUTOSWITCH_PURPLE}rmvenv && mkvenv${AUTOSWITCH_NORMAL} in this directory\n"
98 return
99 fi
100
101 local py_version="$(_python_version "$venv_dir/bin/python")"
102 local message="${AUTOSWITCH_MESSAGE_FORMAT:-"$DEFAULT_MESSAGE_FORMAT"}"
103 message="${message//\%venv_type/$venv_type}"
104 message="${message//\%venv_name/$venv_name}"
105 message="${message//\%py_version/$py_version}"
106 _autoswitch_message "${message}\n"
107
108 # If we are using pipenv and activate its virtual environment - turn down its verbosity
109 # to prevent users seeing " Pipenv found itself running within a virtual environment" warning
110 if [[ "$venv_type" == "pipenv" && "$PIPENV_VERBOSITY" != -1 ]]; then
111 export PIPENV_VERBOSITY=-1
112 fi
113
114 # Much faster to source the activate file directly rather than use the `workon` command
115 local activate_script="$venv_dir/bin/activate"
116
117 _validated_source "$activate_script"
118 fi
119}
120
121
122# Gives the path to the nearest target file
123function _check_path()
124{
125 local check_dir="$1"
126
127 if [[ -f "${check_dir}/${AUTOSWITCH_FILE}" ]]; then
128 printf "${check_dir}/${AUTOSWITCH_FILE}"
129 return
130 elif _is_valid_virtualenv "${check_dir}/${AUTOSWITCH_FILE}"; then
131 printf "${check_dir}/${AUTOSWITCH_FILE}"
132 return
133 elif [[ -f "${check_dir}/poetry.lock" ]]; then
134 printf "${check_dir}/poetry.lock"
135 elif [[ -f "${check_dir}/Pipfile" ]]; then
136 printf "${check_dir}/Pipfile"
137 elif [[ -f "${check_dir}/uv.lock" ]]; then
138 printf "${check_dir}/uv.lock"
139 else
140 # Abort search at file system root or HOME directory (latter is a performance optimisation).
141 if [[ "$check_dir" = "/" || "$check_dir" = "$HOME" ]]; then
142 return
143 fi
144 _check_path "$(dirname "$check_dir")"
145 fi
146}
147
148
149function _activate_poetry() {
150 # check if any environments exist before trying to activate
151 # if env list is empty, then no environment exists that can be activated
152 local name="$(poetry env list --full-path | sort -k 2 | tail -n 1 | cut -d' ' -f1)"
153 if [[ -n "$name" ]]; then
154 _maybeworkon "$name" "poetry"
155 return 0
156 fi
157 return 1
158}
159
160
161function _activate_pipenv() {
162 # unfortunately running pipenv each time we are in a pipenv project directory is slow :(
163 if venv_path="$(PIPENV_IGNORE_VIRTUALENVS=1 pipenv --venv 2>/dev/null)"; then
164 _maybeworkon "$venv_path" "pipenv"
165 return 0
166 fi
167 return 1
168}
169
170function _activate_uv() {
171 if [[ -d ".venv" ]]; then
172 _maybeworkon ".venv" "uv"
173 return 0
174 fi
175 return 1
176}
177
178
179# Automatically switch virtualenv when $AUTOSWITCH_FILE file detected
180function check_venv()
181{
182 local file_owner
183 local file_permissions
184
185 # Get the $AUTOSWITCH_FILE, scanning parent directories
186 local venv_path="$(_check_path "$PWD")"
187
188 if [[ -n "$venv_path" ]]; then
189
190 /usr/bin/stat --version &> /dev/null
191 if [[ $? -eq 0 ]]; then # Linux, or GNU stat
192 file_owner="$(/usr/bin/stat -c %u "$venv_path")"
193 file_permissions="$(/usr/bin/stat -c %a "$venv_path")"
194 else # macOS, or FreeBSD stat
195 file_owner="$(/usr/bin/stat -f %u "$venv_path")"
196 file_permissions="$(/usr/bin/stat -f %OLp "$venv_path")"
197 fi
198
199 if [[ -f "$venv_path" ]] && [[ "$file_owner" != "$(id -u)" ]]; then
200 printf "AUTOSWITCH WARNING: Virtualenv will not be activated\n\n"
201 printf "Reason: Found a $AUTOSWITCH_FILE file but it is not owned by the current user\n"
202 printf "Change ownership of ${AUTOSWITCH_PURPLE}$venv_path${AUTOSWITCH_NORMAL} to ${PURPLE}'$USER'${NORMAL} to fix this\n"
203 elif [[ -f "$venv_path" ]] && ! [[ "$file_permissions" =~ ^[64][04][04]$ ]]; then
204 printf "AUTOSWITCH WARNING: Virtualenv will not be activated\n\n"
205 printf "Reason: Found a $AUTOSWITCH_FILE file with weak permission settings ($file_permissions).\n"
206 printf "Run the following command to fix this: ${AUTOSWITCH_PURPLE}\"chmod 600 $venv_path\"${AUTOSWITCH_NORMAL}\n"
207 else
208 if [[ "$venv_path" == *"/Pipfile" ]]; then
209 if type "pipenv" > /dev/null && _activate_pipenv; then
210 return
211 fi
212 elif [[ "$venv_path" == *"/poetry.lock" ]]; then
213 if type "poetry" > /dev/null && _activate_poetry; then
214 return
215 fi
216 elif [[ "$venv_path" == *"/uv.lock" ]]; then
217 local uv_venv_path="$(dirname "$venv_path")/.venv"
218 _maybeworkon "$uv_venv_path" "uv"
219 return
220 # standard use case: $venv_path is a file containing a virtualenv name
221 elif [[ -f "$venv_path" ]]; then
222 local switch_to="$(<"$venv_path")"
223 _maybeworkon "$(_virtual_env_dir "$switch_to")" "virtualenv"
224 return
225 elif _is_valid_virtualenv "$venv_path"; then
226 _maybeworkon "$venv_path" "virtualenv"
227 return
228 fi
229 fi
230 fi
231
232 _default_venv
233}
234
235
236# Switch to the default virtual environment
237function _default_venv()
238{
239 local venv_type="$(_get_venv_type "$OLDPWD")"
240 if [[ -n "$AUTOSWITCH_DEFAULTENV" ]]; then
241 _maybeworkon "$(_virtual_env_dir "$AUTOSWITCH_DEFAULTENV")" "$venv_type"
242 elif [[ -n "$VIRTUAL_ENV" ]]; then
243 local venv_name="$(_get_venv_name "$VIRTUAL_ENV" "$venv_type")"
244 _autoswitch_message "Deactivating: ${AUTOSWITCH_BOLD}${AUTOSWITCH_PURPLE}%s${AUTOSWITCH_NORMAL}\n" "$venv_name"
245 deactivate
246 fi
247}
248
249
250# remove project environment for current directory
251function rmvenv()
252{
253 local venv_type="$(_get_venv_type "$PWD" "unknown")"
254
255 if [[ "$venv_type" == "pipenv" ]]; then
256 deactivate
257 pipenv --rm
258 elif [[ "$venv_type" == "poetry" ]]; then
259 deactivate
260 poetry env remove "$(poetry run which python)"
261 elif [[ "$venv_type" == "uv" ]]; then
262 deactivate
263 rm -rf ".venv"
264 else
265 if [[ -f "$AUTOSWITCH_FILE" ]]; then
266 local venv_name="$(<$AUTOSWITCH_FILE)"
267
268 # detect if we need to switch virtualenv first
269 if [[ -n "$VIRTUAL_ENV" ]]; then
270 local current_venv="$(basename $VIRTUAL_ENV)"
271 if [[ "$current_venv" = "$venv_name" ]]; then
272 _default_venv
273 fi
274 fi
275
276 printf "Removing ${AUTOSWITCH_PURPLE}%s${AUTOSWITCH_NORMAL}...\n" "$venv_name"
277 # Using explicit paths to avoid any alias/function interference.
278 # rm should always be found in this location according to
279 # https://refspecs.linuxfoundation.org/FHS_3.0/fhs/ch03s04.html
280 # https://www.freedesktop.org/wiki/Software/systemd/TheCaseForTheUsrMerge/
281 /bin/rm -rf "$(_virtual_env_dir "$venv_name")"
282 if [[ -z "$(/bin/ls -A $VIRTUAL_ENV_DIR)" ]]; then
283 /bin/rm -rf "$VIRTUAL_ENV_DIR"
284 fi
285 /bin/rm "$AUTOSWITCH_FILE"
286 else
287 printf "No $AUTOSWITCH_FILE file in the current directory!\n"
288 fi
289 fi
290}
291
292
293function _missing_error_message() {
294 local command="$1"
295 printf "${AUTOSWITCH_BOLD}${AUTOSWITCH_RED}"
296 printf "zsh-autoswitch-virtualenv requires '%s' to install this project!\n\n" "$command"
297 printf "${AUTOSWITCH_NORMAL}"
298 printf "If this is already installed but you are still seeing this message, \n"
299 printf "then make sure the ${AUTOSWITCH_BOLD}$command${AUTOSWITCH_NORMAL} command is in your PATH.\n" $command
300 printf "\n"
301}
302
303function randstr()
304{
305 local python_bin="${1:-${AUTOSWITCH_DEFAULT_PYTHON:-python3}}"
306 "$python_bin" -c "from __future__ import print_function; import string, random; print(''.join(random.choice(string.ascii_lowercase) for _ in range(4)))"
307}
308
309
310# helper function to create a project environment for the current directory
311function mkvenv()
312{
313 local venv_type="$(_get_venv_type "$PWD" "unknown")"
314 # Copy parameters variable so that we can mutate it
315 # NOTE: Keep declaration of variable and assignment separate for zsh 5.0 compatibility
316 local params
317 params=("${@[@]}")
318
319 if [[ "$venv_type" == "pipenv" ]]; then
320 if ! type "pipenv" > /dev/null; then
321 _missing_error_message pipenv
322 return
323 fi
324 # TODO: detect if this is already installed
325 if [[ "$AUTOSWITCH_PIPINSTALL" = "FULL" ]]
326 then
327 pipenv install --dev $params
328 else
329 pipenv install --dev --editable . $params
330 fi
331 _activate_pipenv
332 return
333 elif [[ "$venv_type" == "poetry" ]]; then
334 if ! type "poetry" > /dev/null; then
335 _missing_error_message poetry
336 return
337 fi
338 # TODO: detect if this is already installed
339 poetry install $params
340 _activate_poetry
341 return
342 elif [[ "$venv_type" == "uv" ]]; then
343 if ! type "uv" > /dev/null; then
344 _missing_error_message uv
345 return
346 fi
347 uv sync $params
348 _activate_uv
349 return
350 else
351 local python_bin="${AUTOSWITCH_DEFAULT_PYTHON:-python3}"
352 local use_default_python=1
353 local verbose=0
354 local -a venv_params
355 local index
356
357 for (( index = 1; index <= ${#params}; index++ )); do
358 if [[ "${params[index]}" == "--python" ]]; then
359 (( index++ ))
360 if (( index > ${#params} )); then
361 printf "mkvenv: --python requires a value\n"
362 return 1
363 fi
364 python_bin="${params[index]}"
365 use_default_python=0
366 elif [[ "${params[index]}" == --python=* ]]; then
367 python_bin="${params[index]#--python=}"
368 use_default_python=0
369 elif [[ "${params[index]}" == "--verbose" ]]; then
370 verbose=1
371 else
372 venv_params+=("${params[index]}")
373 fi
374 done
375
376 if [[ -z "$python_bin" ]]; then
377 printf "mkvenv: --python requires a value\n"
378 return 1
379 fi
380
381 if ! type "$python_bin" > /dev/null; then
382 _missing_error_message "$python_bin"
383 return
384 fi
385
386 if [[ -f "$AUTOSWITCH_FILE" ]]; then
387 printf "$AUTOSWITCH_FILE file already exists. If this is a mistake use the rmvenv command\n"
388 else
389 local venv_name="$(basename $PWD)-$(randstr "$python_bin")"
390 local venv_path="$(_virtual_env_dir "$venv_name")"
391
392 printf "Creating ${AUTOSWITCH_PURPLE}%s${NONE} virtualenv\n" "$venv_name"
393
394 if (( use_default_python )) && [[ -n "$AUTOSWITCH_DEFAULT_PYTHON" ]]; then
395 printf "${AUTOSWITCH_PURPLE}"
396 printf 'Using $AUTOSWITCH_DEFAULT_PYTHON='
397 printf "$AUTOSWITCH_DEFAULT_PYTHON"
398 printf "${NONE}\n"
399 fi
400
401 /bin/mkdir -p "$VIRTUAL_ENV_DIR"
402
403 if (( verbose )); then
404 "$python_bin" -m venv "${venv_params[@]}" "$venv_path" || return
405 else
406 "$python_bin" -m venv "${venv_params[@]}" "$venv_path" > /dev/null || return
407 fi
408
409 printf "$venv_name\n" > "$AUTOSWITCH_FILE"
410 chmod 600 "$AUTOSWITCH_FILE"
411
412 _maybeworkon "$venv_path" "virtualenv"
413 fi
414 fi
415}
416
417
418function enable_autoswitch_virtualenv() {
419 disable_autoswitch_virtualenv
420 add-zsh-hook chpwd check_venv
421}
422
423
424function disable_autoswitch_virtualenv() {
425 add-zsh-hook -D chpwd check_venv
426}
427
428# This function is only used to startup zsh-autoswitch-virtualenv
429# the first time a terminal is started up
430# it waits for the terminal to be ready using precmd and then
431# immediately removes itself from the zsh-hook.
432# This seems important for "instant prompt" zsh themes like powerlevel10k
433function _autoswitch_startup() {
434 local python_bin="${AUTOSWITCH_DEFAULT_PYTHON:-python3}"
435 if ! type "${python_bin}" > /dev/null; then
436 printf "WARNING: python binary '${python_bin}' not found on PATH.\n"
437 printf "zsh-autoswitch-virtualenv plugin will be disabled.\n"
438 else
439 enable_autoswitch_virtualenv
440 check_venv
441 fi
442 add-zsh-hook -D precmd _autoswitch_startup
443}
444
445autoload -Uz add-zsh-hook
446add-zsh-hook precmd _autoswitch_startup