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