From 72707d92b1f127a2c16f0447e4d8233ed7b3c10e Mon Sep 17 00:00:00 2001
From: 0b10 <qbit8@protonmail.com>
Date: Mon, 2 Mar 2020 13:38:44 +0100
Subject: [PATCH] python: add pyuserpaths function (#7758)

Summary
    Make it get a list of installed interpreters, and add the
    relevant local (HOME) site-packages directory to PYTHONPATH.

    Reason
    To easily add all relevant paths, initially, and between upgrades.

Add check for PYTHONUSERBASE

    Summary
    Check for a non-standard install directory, use it if one
    exists, otherwise use ~/.local.

    Reason
    Allow users to specify their own installation directory, without
    affecting functionality.
---
 plugins/python/README.md         |  3 ++-
 plugins/python/python.plugin.zsh | 29 +++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/plugins/python/README.md b/plugins/python/README.md
index e391bcfd4..0180218a2 100644
--- a/plugins/python/README.md
+++ b/plugins/python/README.md
@@ -11,7 +11,8 @@ plugins=(... python)
 
 | Command          | Description                                                                     |
 |------------------|---------------------------------------------------------------------------------|
+| `ipython`        | Runs the appropriate `ipython` version according to the activated virtualenv    |
 | `pyfind`         | Finds .py files recursively in the current directory                            |
 | `pyclean [dirs]` | Deletes byte-code and cache files from a list of directories or the current one |
 | `pygrep <text>`  | Looks for `text` in .py files                                                   |
-| `ipython`        | Runs the appropriate `ipython` version according to the activated virtualenv    |
+| `pyuserpaths`    | Add --user site-packages to PYTHONPATH, for all installed python versions.      |
diff --git a/plugins/python/python.plugin.zsh b/plugins/python/python.plugin.zsh
index 62855fca2..3e4db1e7d 100644
--- a/plugins/python/python.plugin.zsh
+++ b/plugins/python/python.plugin.zsh
@@ -11,6 +11,35 @@ function pyclean() {
     find ${ZSH_PYCLEAN_PLACES} -depth -type d -name ".pytest_cache" -exec rm -r "{}" +
 }
 
+# Add the user installed site-packages paths to PYTHONPATH, only if the
+#   directory exists. Also preserve the current PYTHONPATH value.
+# Feel free to autorun this when .zshrc loads.
+function pyuserpaths() {
+    local targets=("python2" "python3")  # bins
+    
+    # Get existing interpreters.
+    local interps=()
+    for target in $targets; do
+        [ `command -v $target` ] && interps+=($target)
+    done
+
+    # Check for a non-standard install directory.
+    local user_base="${HOME}/.local"
+    [ $PYTHONUSERBASE ] && user_base=$PYTHONUSERBASE
+
+    # Add version specific paths, if:
+    #   it exists in the filesystem;
+    #   it isn't in PYTHONPATH already.
+    for interp in $interps; do
+        # Get minor release version.
+        local ver=`$interp -V 2>&1`
+        ver=`echo ${ver:7} | cut -d '.' -f 1,2`  # The patch version is variable length, truncate it.
+        
+        local site_pkgs="${user_base}/lib/python${ver}/site-packages"
+        [[ -d $site_pkgs && ! $PYTHONPATH =~ $site_pkgs ]] && export PYTHONPATH=${site_pkgs}:$PYTHONPATH
+    done
+}
+
 # Grep among .py files
 alias pygrep='grep -r --include="*.py"'