make runtime env toolchain work with bootstrap=script venv

Author: rickeylev

python/private/py_executable.bzl

@@ -350,6 +350,7 @@ def _create_executable(
             main_py = main_py,
             imports = imports,
             runtime_details = runtime_details,
+            venv = venv,
         )
         extra_runfiles = ctx.runfiles([stage2_bootstrap] + venv.files_without_interpreter)
         zip_main = _create_zip_main(
@@ -557,6 +558,8 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         ctx.actions.write(interpreter, "actual:{}".format(interpreter_actual_path))
 
     elif runtime.interpreter:
+        # Some wrappers around the interpreter (e.g. pyenv) use the program
+        # name to decide what to do, so preserve the name.
         py_exe_basename = paths.basename(runtime.interpreter.short_path)
 
         # Even though ctx.actions.symlink() is used, using
@@ -594,7 +597,8 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
     if "t" in runtime.abi_flags:
         version += "t"
 
-    site_packages = "{}/lib/python{}/site-packages".format(venv, version)
+    venv_site_packages = "lib/python{}/site-packages".format(version)
+    site_packages = "{}/{}".format(venv, venv_site_packages)
     pth = ctx.actions.declare_file("{}/bazel.pth".format(site_packages))
     ctx.actions.write(pth, "import _bazel_site_init\n")
 
@@ -620,6 +624,7 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         # Runfiles root relative path or absolute path
         interpreter_actual_path = interpreter_actual_path,
         files_without_interpreter = [pyvenv_cfg, pth, site_init] + site_packages_symlinks,
+        venv_site_packages = venv_site_packages,
     )
 
 def _create_site_packages_symlinks(ctx, site_packages):
@@ -716,7 +721,8 @@ def _create_stage2_bootstrap(
         output_sibling,
         main_py,
         imports,
-        runtime_details):
+        runtime_details,
+        venv = None):
     output = ctx.actions.declare_file(
         # Prepend with underscore to prevent pytest from trying to
         # process the bootstrap for files starting with `test_`
@@ -742,6 +748,7 @@ def _create_stage2_bootstrap(
             "%main_module%": ctx.attr.main_module,
             "%target%": str(ctx.label),
             "%workspace_name%": ctx.workspace_name,
+            "%site_packages%": venv.venv_site_packages if venv else "",
         },
         is_executable = True,
     )

python/private/runtime_env_toolchain_interpreter.sh

@@ -53,5 +53,27 @@ documentation for py_runtime_pair \
 (https://github.com/bazel-contrib/rules_python/blob/master/docs/python.md#py_runtime_pair)."
 fi
 
-exec "$PYTHON_BIN" "$@"
+# Because this is a wrapper script that invokes Python, it prevents Python from
+# detecting virtualenvs like normal (i.e. using the venv symlink to find the
+# real interpreter). To work around this, we have to manually detect the venv,
+# then trick the interpreter into understanding we're in a virtual env.
+self_dir=$(dirname $0)
+if [[ -e "$self_dir/pyvenv.cfg" || -e "$self_dir/../pyvenv.cfg" ]]; then
+  if [[ "$0" == /* ]]; then
+    venv_bin="$0"
+  else
+    venv_bin="$PWD/$0"
+  fi
 
+  # PYTHONEXECUTABLE is also used because `exec -a` doesn't fully trick the
+  # pyenv wrappers.
+  # NOTE: The PYTHONEXECUTABLE envvar docs say it's only for Mac, however,
+  # emperically it also works on e.g. Linux
+  export PYTHONEXECUTABLE="$venv_bin"
+  # Python uses to det
+  # Use exec -a so that Python looks at the venv's binary, not the
+  # actual one invoked.
+  exec -a "$venv_bin" "$PYTHON_BIN" "$@"
+else
+  exec "$PYTHON_BIN" "$@"
+fi

python/private/stage2_bootstrap_template.py

@@ -32,6 +32,10 @@
 # Module name to execute. Empty if MAIN is used.
 MAIN_MODULE = "%main_module%"
 
+# venv-relative path to the expected location of the binary's site-packages
+# directory.
+VENV_SITE_PACKAGES = "%site_packages%"
+
 # ===== Template substitutions end =====
 
 
@@ -365,6 +369,18 @@ def main():
     print_verbose("initial environ:", mapping=os.environ)
     print_verbose("initial sys.path:", values=sys.path)
 
+    site_packages = os.path.join(sys.prefix, VENV_SITE_PACKAGES)
+    if site_packages not in sys.path:
+        # NOTE: if this happens, it likely means we're running with a different
+        # Python version than was built with. Things may or may not work.
+        # Such a situation is likely due to the runtime_env toolchain, or some
+        # toolchain configuration. In any case, this better matches how the
+        # previous bootstrap=system_python bootstrap worked (using PYTHONPATH,
+        # which isn't version-specific).
+        print_verbose(f"sys.path missing expected site-packages: adding {site_packages}")
+        import site
+        site.addsitedir(site_packages)
+
     main_rel_path = None
     # todo: things happen to work because find_runfiles_root
     # ends up using stage2_bootstrap, and ends up computing the proper