From 819e86b69f3a01ee3cc5518078242be3a2eb2854 Mon Sep 17 00:00:00 2001
From: GFJ138 <sebastien.dementen@engie.com>
Date: Mon, 1 Aug 2022 15:10:04 +0200
Subject: [PATCH] add the --pip-use-shebang-python option

---
 docs/cli-reference.rst |  4 ++++
 src/shiv/cli.py        | 35 +++++++++++++++++++++++++++++++++--
 src/shiv/pip.py        |  6 +++---
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst
index c80236a..69e7d4c 100644
--- a/docs/cli-reference.rst
+++ b/docs/cli-reference.rst
@@ -40,3 +40,7 @@ Also note that you can always fix the shebang during installation of a zipapp us
 .. code-block:: shell
 
    python3 -m zipapp -p '/usr/bin/env python3.7' -o ~/bin/foo foo.pyz
+
+If the interpreter you run is different from the interpreter of the sheband line,
+you can use the :option:`--pip-use-shebang-python` to use the shebang interpreter when installing packages through pip
+(important for version specific packages).
\ No newline at end of file
diff --git a/src/shiv/cli.py b/src/shiv/cli.py
index e461feb..3c76083 100644
--- a/src/shiv/cli.py
+++ b/src/shiv/cli.py
@@ -1,6 +1,7 @@
 import hashlib
 import os
 import shutil
+import subprocess
 import sys
 import time
 
@@ -156,6 +157,11 @@ def copytree(src: Path, dst: Path) -> None:
     ),
 )
 @click.option("--root", type=click.Path(), help="Override the 'root' path (default is ~/.shiv).")
+@click.option(
+    "--pip-use-shebang-python",
+    is_flag=True,
+    help="Use the python interpreter from the shebang line when installing the pip packages.",
+)
 @click.argument("pip_args", nargs=-1, type=click.UNPROCESSED)
 def main(
     output_file: str,
@@ -172,6 +178,7 @@ def main(
     preamble: Optional[str],
     root: Optional[str],
     pip_args: List[str],
+    pip_use_shebang_python: bool,
 ) -> None:
     """
     Shiv is a command line utility for building fully self-contained Python zipapps
@@ -181,6 +188,14 @@ def main(
     if not pip_args and not site_packages:
         sys.exit(NO_PIP_ARGS_OR_SITE_PACKAGES)
 
+    if python:
+        python = str(python)
+    else:
+        if os.name == "nt":
+            python = sys.executable
+        else:
+            python = DEFAULT_SHEBANG
+
     if output_file is None:
         sys.exit(NO_OUTFILE)
 
@@ -212,7 +227,23 @@ def main(
 
         if pip_args:
             # Install dependencies into staged site-packages.
-            pip.install(["--target", tmp_site_packages] + list(pip_args))
+            if pip_use_shebang_python:
+                assert python is not None, "You should have specified a python interpreter."
+                pip_interpreter = python
+            else:
+                version_sys = subprocess.run([sys.executable, "-V"], capture_output=True).stdout.decode().rstrip()
+                version_shebang = subprocess.run([python, "-V"], capture_output=True).stdout.decode().rstrip()
+                if version_sys != version_shebang:
+                    click.secho(
+                        f"You are installing pip packages with your current python interpreter ({version_sys}) "
+                        f"while the python interpreter used in your shebang line is ({version_shebang}). "
+                        f"If you want to use the python interpreter from your shebang line to install the "
+                        f"pip packages, you must activate the option 'pip_use_shebang_python'.",
+                        err=True,
+                        fg="red",
+                    )
+                pip_interpreter = sys.executable
+            pip.install(["--target", tmp_site_packages] + list(pip_args), pip_interpreter=pip_interpreter)
 
         if preamble:
             bin_dir = Path(tmp_site_packages, "bin")
@@ -269,7 +300,7 @@ def main(
         builder.create_archive(
             sources,
             target=Path(output_file).expanduser(),
-            interpreter=python or DEFAULT_SHEBANG,
+            interpreter=python,
             main="_bootstrap:bootstrap",
             env=env,
             compressed=compressed,
diff --git a/src/shiv/pip.py b/src/shiv/pip.py
index 4457577..cab2fa1 100644
--- a/src/shiv/pip.py
+++ b/src/shiv/pip.py
@@ -28,10 +28,10 @@ def clean_pip_env() -> Generator[None, None, None]:
             os.environ[PIP_REQUIRE_VIRTUALENV] = require_venv
 
 
-def install(args: List[str]) -> None:
+def install(args: List[str], pip_interpreter: str) -> None:
     """`pip install` as a function.
 
-    Accepts a list of pip arguments.
+    Accepts a list of pip arguments and the python interpreter to use when running pip.
 
     .. code-block:: py
 
@@ -54,7 +54,7 @@ def install(args: List[str]) -> None:
         extend_python_path(subprocess_env, sys.path[sitedir_index:])
 
         process = subprocess.Popen(
-            [sys.executable, "-m", "pip", "--disable-pip-version-check", "install", *args],
+            [pip_interpreter, "-m", "pip", "--disable-pip-version-check", "install", *args],
             stdout=subprocess.PIPE,
             stderr=subprocess.STDOUT,
             env=subprocess_env,