wip: make sys.executable work with script bootstrap

Author: rickeylev

CHANGELOG.md

@@ -52,15 +52,20 @@ Unreleased changes template.
 
 {#v0-0-0-changed}
 ### Changed
-* Nothing yet.
+* (binaries/tests) For {obj}`--bootstrap_impl=script`, an empty,
+  binary-specific, virtual env is used to customize sys.path initialization.
 
 {#v0-0-0-fixed}
 ### Fixed
-* Nothing yet.
+* (binaries/tests) ({obj}`--bootstrap_impl=scipt`) Using `sys.executable` will
+  use the same `sys.path` setup as the calling binary.
+  ([2169](https://github.com/bazelbuild/rules_python/issues/2169)).
 
 {#v0-0-0-added}
 ### Added
-* Nothing yet.
+* (providers) Added {obj}`py_runtime_info.site_init_template` and
+  {obj}`PyRuntimeInfo.site_init_template` for specifying the template to use to
+  initialize the interpreter via venv startup hooks.
 
 {#v0-0-0-removed}
 ### Removed
@@ -99,6 +104,9 @@ Unreleased changes template.
 * (precompiling) Skip precompiling (instead of erroring) if the legacy
   `@bazel_tools//tools/python:autodetecting_toolchain` is being used
   ([#2364](https://github.com/bazelbuild/rules_python/issues/2364)).
+* (bzlmod) Generate `config_setting` values for all available toolchains instead
+  of only the registered toolchains, which restores the previous behaviour that
+  `bzlmod` users would have observed.
 
 {#v0-39-0-added}
 ### Added

python/private/BUILD.bazel

@@ -702,6 +702,14 @@ filegroup(
     visibility = ["//visibility:public"],
 )
 
+filegroup(
+    name = "site_init_template",
+    srcs = ["site_init_template.py"],
+    # Not actually public. Only public because it's an implicit dependency of
+    # py_runtime.
+    visibility = ["//visibility:public"],
+)
+
 # NOTE: Windows builds don't use this bootstrap. Instead, a native Windows
 # program locates some Python exe and runs `python.exe foo.zip` which
 # runs the __main__.py in the zip file.

python/private/py_executable_bazel.bzl

@@ -81,6 +81,9 @@ the `srcs` of Python targets as required.
         "_py_toolchain_type": attr.label(
             default = TARGET_TOOLCHAIN_TYPE,
         ),
+        "_python_version_flag": attr.label(
+            default = "//python/config_settings:python_version",
+        ),
         "_windows_launcher_maker": attr.label(
             default = "@bazel_tools//tools/launcher:launcher_maker",
             cfg = "exec",
@@ -177,13 +180,22 @@ def _create_executable(
     else:
         base_executable_name = executable.basename
 
+    venv = None
+
     # The check for stage2_bootstrap_template is to support legacy
     # BuiltinPyRuntimeInfo providers, which is likely to come from
     # @bazel_tools//tools/python:autodetecting_toolchain, the toolchain used
     # for workspace builds when no rules_python toolchain is configured.
     if (BootstrapImplFlag.get_value(ctx) == BootstrapImplFlag.SCRIPT and
         runtime_details.effective_runtime and
         hasattr(runtime_details.effective_runtime, "stage2_bootstrap_template")):
+        venv = _create_venv(
+            ctx,
+            output_prefix = base_executable_name,
+            imports = imports,
+            runtime_details = runtime_details,
+        )
+
         stage2_bootstrap = _create_stage2_bootstrap(
             ctx,
             output_prefix = base_executable_name,
@@ -192,11 +204,12 @@ def _create_executable(
             imports = imports,
             runtime_details = runtime_details,
         )
-        extra_runfiles = ctx.runfiles([stage2_bootstrap])
+        extra_runfiles = ctx.runfiles([stage2_bootstrap] + venv.files_without_interpreter)
         zip_main = _create_zip_main(
             ctx,
             stage2_bootstrap = stage2_bootstrap,
             runtime_details = runtime_details,
+            venv = venv,
         )
     else:
         stage2_bootstrap = None
@@ -272,6 +285,7 @@ def _create_executable(
             zip_file = zip_file,
             stage2_bootstrap = stage2_bootstrap,
             runtime_details = runtime_details,
+            venv = venv,
         )
     elif bootstrap_output:
         _create_stage1_bootstrap(
@@ -282,6 +296,7 @@ def _create_executable(
             is_for_zip = False,
             imports = imports,
             main_py = main_py,
+            venv = venv,
         )
     else:
         # Otherwise, this should be the Windows case of launcher + zip.
@@ -296,13 +311,20 @@ def _create_executable(
                 build_zip_enabled = build_zip_enabled,
             ))
 
+    # The interpreter is added this late in the process so that it isn't
+    # added to the files that zipping processes.
+    if venv:
+        extra_runfiles = extra_runfiles.merge(ctx.runfiles([venv.interpreter]))
     return create_executable_result_struct(
         extra_files_to_build = depset(extra_files_to_build),
         output_groups = {"python_zip_file": depset([zip_file])},
         extra_runfiles = extra_runfiles,
     )
 
-def _create_zip_main(ctx, *, stage2_bootstrap, runtime_details):
+def _create_zip_main(ctx, *, stage2_bootstrap, runtime_details, venv):
+    python_binary = _runfiles_root_path(ctx, venv.interpreter.short_path)
+    python_binary_actual = _runfiles_root_path(ctx, venv.interpreter_actual_path)
+
     # The location of this file doesn't really matter. It's added to
     # the zip file as the top-level __main__.py file and not included
     # elsewhere.
@@ -311,7 +333,8 @@ def _create_zip_main(ctx, *, stage2_bootstrap, runtime_details):
         template = runtime_details.effective_runtime.zip_main_template,
         output = output,
         substitutions = {
-            "%python_binary%": runtime_details.executable_interpreter_path,
+            "%python_binary%": python_binary,
+            "%python_binary_actual%": python_binary_actual,
             "%stage2_bootstrap%": "{}/{}".format(
                 ctx.workspace_name,
                 stage2_bootstrap.short_path,
@@ -321,6 +344,75 @@ def _create_zip_main(ctx, *, stage2_bootstrap, runtime_details):
     )
     return output
 
+# Create a venv the executable can use.
+# For venv details and the venv startup process, see:
+# * https://docs.python.org/3/library/venv.html
+# * https://snarky.ca/how-virtual-environments-work/
+# * https://github.com/python/cpython/blob/main/Modules/getpath.py
+# * https://github.com/python/cpython/blob/main/Lib/site.py
+def _create_venv(ctx, output_prefix, imports, runtime_details):
+    venv = "_{}.venv".format(output_prefix.lstrip("_"))
+
+    # The pyvenv.cfg file must be present to trigger the venv site hooks.
+    # Because it's paths are expected to be absolute paths, we can't reliably
+    # put much in it. See https://github.com/python/cpython/issues/83650
+    pyvenv_cfg = ctx.actions.declare_file("{}/pyvenv.cfg".format(venv))
+    ctx.actions.write(pyvenv_cfg, "")
+
+    runtime = runtime_details.effective_runtime
+    if runtime.interpreter:
+        py_exe_basename = paths.basename(runtime.interpreter.short_path)
+        interpreter = ctx.actions.declare_file("{}/bin/{}".format(venv, py_exe_basename))
+        ctx.actions.symlink(output = interpreter, target_file = runtime.interpreter)
+        interpreter_actual_path = runtime.interpreter.short_path
+    else:
+        py_exe_basename = paths.basename(runtime.interpreter_path)
+        interpreter = ctx.actions.declare_symlink("{}/bin/{}".format(venv, py_exe_basename))
+        ctx.actions.symlink(output = interpreter, target_path = runtime.interpreter_path)
+        interpreter_actual_path = runtime.interpreter_path
+
+    if runtime.interpreter_version_info:
+        version = "{}.{}".format(
+            runtime.interpreter_version_info.major,
+            runtime.interpreter_version_info.minor,
+        )
+    else:
+        version_flag = ctx.attr._python_version_flag[config_common.FeatureFlagInfo].value
+        version_flag_parts = version_flag.split(".")[0:2]
+        version = "{}.{}".format(*version_flag_parts)
+
+    # See site.py logic: free-threaded builds append "t" to the venv lib dir name
+    if "t" in runtime.abi_flags:
+        version += "t"
+
+    site_packages = "{}/lib/python{}/site-packages".format(venv, version)
+    pth = ctx.actions.declare_file("{}/bazel.pth".format(site_packages))
+    ctx.actions.write(pth, "import _bazel_site_init\n")
+
+    site_init = ctx.actions.declare_file("{}/_bazel_site_init.py".format(site_packages))
+    computed_subs = ctx.actions.template_dict()
+    computed_subs.add_joined("%imports%", imports, join_with = ":", map_each = _map_each_identity)
+    ctx.actions.expand_template(
+        template = runtime.site_init_template,
+        output = site_init,
+        substitutions = {
+            "%import_all%": "True" if ctx.fragments.bazel_py.python_import_all_repositories else "False",
+            "%site_init_runfiles_path%": "{}/{}".format(ctx.workspace_name, site_init.short_path),
+            "%workspace_name%": ctx.workspace_name,
+        },
+        computed_substitutions = computed_subs,
+    )
+
+    return struct(
+        interpreter = interpreter,
+        # Runfiles-relative path or absolute path
+        interpreter_actual_path = interpreter_actual_path,
+        files_without_interpreter = [pyvenv_cfg, pth, site_init],
+    )
+
+def _map_each_identity(v):
+    return v
+
 def _create_stage2_bootstrap(
         ctx,
         *,
@@ -363,6 +455,13 @@ def _create_stage2_bootstrap(
     )
     return output
 
+def _runfiles_root_path(ctx, path):
+    # The ../ comes from short_path for files in other repos.
+    if path.startswith("../"):
+        return path[3:]
+    else:
+        return "{}/{}".format(ctx.workspace_name, path)
+
 def _create_stage1_bootstrap(
         ctx,
         *,
@@ -371,12 +470,24 @@ def _create_stage1_bootstrap(
         stage2_bootstrap = None,
         imports = None,
         is_for_zip,
-        runtime_details):
+        runtime_details,
+        venv = None):
     runtime = runtime_details.effective_runtime
 
+    if venv:
+        python_binary_path = _runfiles_root_path(ctx, venv.interpreter.short_path)
+    else:
+        python_binary_path = runtime_details.executable_interpreter_path
+
+    if is_for_zip and venv:
+        python_binary_actual = _runfiles_root_path(ctx, venv.interpreter_actual_path)
+    else:
+        python_binary_actual = ""
+
     subs = {
         "%is_zipfile%": "1" if is_for_zip else "0",
-        "%python_binary%": runtime_details.executable_interpreter_path,
+        "%python_binary%": python_binary_path,
+        "%python_binary_actual%": python_binary_actual,
         "%target%": str(ctx.label),
         "%workspace_name%": ctx.workspace_name,
     }
@@ -447,6 +558,7 @@ def _create_windows_exe_launcher(
     )
 
 def _create_zip_file(ctx, *, output, original_nonzip_executable, zip_main, runfiles):
+    """Create a Python zipapp (zip with __main__.py entry point)."""
     workspace_name = ctx.workspace_name
     legacy_external_runfiles = _py_builtins.get_legacy_external_runfiles(ctx)
 
@@ -524,7 +636,14 @@ def _get_zip_runfiles_path(path, workspace_name, legacy_external_runfiles):
         zip_runfiles_path = paths.normalize("{}/{}".format(workspace_name, path))
     return "{}/{}".format(_ZIP_RUNFILES_DIRECTORY_NAME, zip_runfiles_path)
 
-def _create_executable_zip_file(ctx, *, output, zip_file, stage2_bootstrap, runtime_details):
+def _create_executable_zip_file(
+        ctx,
+        *,
+        output,
+        zip_file,
+        stage2_bootstrap,
+        runtime_details,
+        venv):
     prelude = ctx.actions.declare_file(
         "{}_zip_prelude.sh".format(output.basename),
         sibling = output,
@@ -536,6 +655,7 @@ def _create_executable_zip_file(ctx, *, output, zip_file, stage2_bootstrap, runt
             stage2_bootstrap = stage2_bootstrap,
             runtime_details = runtime_details,
             is_for_zip = True,
+            venv = venv,
         )
     else:
         ctx.actions.write(prelude, "#!/usr/bin/env python3\n")

python/private/py_runtime_info.bzl

@@ -68,7 +68,8 @@ def _PyRuntimeInfo_init(
         interpreter_version_info = None,
         stage2_bootstrap_template = None,
         zip_main_template = None,
-        abi_flags = ""):
+        abi_flags = "",
+        site_init_template = None):
     if (interpreter_path and interpreter) or (not interpreter_path and not interpreter):
         fail("exactly one of interpreter or interpreter_path must be specified")
 
@@ -117,6 +118,7 @@ def _PyRuntimeInfo_init(
         "interpreter_version_info": interpreter_version_info_struct_from_dict(interpreter_version_info),
         "pyc_tag": pyc_tag,
         "python_version": python_version,
+        "site_init_template": site_init_template,
         "stage2_bootstrap_template": stage2_bootstrap_template,
         "stub_shebang": stub_shebang,
         "zip_main_template": zip_main_template,
@@ -126,6 +128,11 @@ PyRuntimeInfo, _unused_raw_py_runtime_info_ctor = define_bazel_6_provider(
     doc = """Contains information about a Python runtime, as returned by the `py_runtime`
 rule.
 
+:::{warning}
+This is an **unstable public** API. It may change more frequently and has weaker
+compatibility guarantees.
+:::
+
 A Python runtime describes either a *platform runtime* or an *in-build runtime*.
 A platform runtime accesses a system-installed interpreter at a known path,
 whereas an in-build runtime points to a `File` that acts as the interpreter. In
@@ -139,6 +146,9 @@ the same conventions as the standard CPython interpreter.
 :type: str
 
 The runtime's ABI flags, i.e. `sys.abiflags`.
+
+:::{versionadded} 0.39.0
+:::
 """,
         "bootstrap_template": """
 :type: File
@@ -160,7 +170,8 @@ is expected to behave and the substutitions performed.
   `%target%`, `%workspace_name`, `%coverage_tool%`, `%import_all%`, `%imports%`,
   `%main%`, `%shebang%`
 * `--bootstrap_impl=script` substititions: `%is_zipfile%`, `%python_binary%`,
-  `%target%`, `%workspace_name`, `%shebang%, `%stage2_bootstrap%`
+  `%python_binary_actual%`, `%target%`, `%workspace_name`,
+  `%shebang%`, `%stage2_bootstrap%`
 
 Substitution definitions:
 
@@ -172,6 +183,19 @@ Substitution definitions:
   * An absolute path to a system interpreter (e.g. begins with `/`).
   * A runfiles-relative path to an interpreter (e.g. `somerepo/bin/python3`)
   * A program to search for on PATH, i.e. a word without spaces, e.g. `python3`.
+
+  When `--bootstrap_impl=script` is used, this is always a runfiles-relative
+  path to a venv-based interpreter executable.
+
+* `%python_binary_actual%`: The path to the interpreter that
+  `%python_binary%` invokes. There are three types of paths:
+  * An absolute path to a system interpreter (e.g. begins with `/`).
+  * A runfiles-relative path to an interpreter (e.g. `somerepo/bin/python3`)
+  * A program to search for on PATH, i.e. a word without spaces, e.g. `python3`.
+
+  Only set for zip builds with `--bootstrap_impl=script`; other builds will use
+  an empty string.
+
 * `%workspace_name%`: The name of the workspace the target belongs to.
 * `%is_zipfile%`: The string `1` if this template is prepended to a zipfile to
   create a self-executable zip file. The string `0` otherwise.
@@ -250,6 +274,15 @@ correctly.
 
 Indicates whether this runtime uses Python major version 2 or 3. Valid values
 are (only) `"PY2"` and `"PY3"`.
+""",
+        "site_init_template": """
+:type: File
+
+The template to use for the binary-specific site-init hook run by the
+interpreter at startup.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
 """,
         "stage2_bootstrap_template": """
 :type: File