feat(venv): make wheel scripts runnable in venv (bazel-contrib#3743)

Currently, `#!python` based scripts in a wheel are put into the venv
as-is, so aren't actually runnable. Additionally, wheel entry points
aren't installed into the venv.

To fix, rewrite the scripts to re-exec themselves with the `python3`
binary in the same directory.

Also adds support for wheel entry points in venvs. This is done by
having the repo-phase parse entry_points.txt and defining a target for
each entry point. These targets run during build phase, so are able
to generate platform-specific outputs.

Cross-building with Windows is also added, supporting both Windows
as a target or exec platform to generate outputs for the target
platform.

Fixes bazel-contrib#3202

Author: rickeylev

examples/pip_parse/pip_parse_test.py

@@ -50,18 +50,22 @@ def test_entry_point(self):
     def test_data(self):
         actual = os.environ.get("WHEEL_DATA_CONTENTS")
         self.assertIsNotNone(actual)
-        actual = self._remove_leading_dirs(actual.split(" "))
+        actual = set(self._remove_leading_dirs(actual.split(" ")))
+
+        s3cmd_bin = "bin/s3cmd"
+        if os.name == "nt":
+            s3cmd_bin += ".bat"
 
-        expected = [
-            "bin/s3cmd",
+        expected = {
+            s3cmd_bin,
             "data/share/doc/packages/s3cmd/INSTALL.md",
             "data/share/doc/packages/s3cmd/LICENSE",
             "data/share/doc/packages/s3cmd/NEWS",
             "data/share/doc/packages/s3cmd/README.md",
             "data/share/man/man1/s3cmd.1",
-        ]
+        }
 
-        self.assertListEqual(actual, expected)
+        self.assertEqual(actual, expected)
 
     def test_dist_info(self):
         actual = os.environ.get("WHEEL_DIST_INFO_CONTENTS")

python/private/py_executable.bzl

@@ -749,6 +749,14 @@ def _create_venv_windows(ctx, *, venv_ctx_rel_root, runtime, interpreter_actual_
             link_to_path = interpreter_actual_path,
             files = depset([runtime.interpreter]),
         ))
+
+        # This isn't strictly correct, but should work ok.
+        interpreter_symlinks.add(ExplicitSymlink(
+            runfiles_path = paths.join(paths.dirname(rf_path), "pythonw.exe"),
+            venv_path = paths.join(paths.dirname(venv_rel_path), "pythonw.exe"),
+            link_to_path = paths.join(paths.dirname(interpreter_actual_path), "pythonw.exe"),
+            files = depset(),
+        ))
     else:
         # It's OK to use declare_symlink here because an absolute path
         # will be written to it, so Bazel won't mangle it.

python/private/pypi/BUILD.bazel

@@ -24,6 +24,24 @@ exports_files(
     visibility = ["//visibility:public"],
 )
 
+alias(
+    name = "venv_entry_point_template",
+    actual = select({
+        "@platforms//os:windows": "venv_entry_point_template.bat",
+        "//conditions:default": "venv_entry_point_template.sh",
+    }),
+    visibility = ["//visibility:public"],
+)
+
+alias(
+    name = "venv_shebang_rewriter",
+    actual = select({
+        "@platforms//os:windows": "venv_shebang_rewriter.ps1",
+        "//conditions:default": "venv_shebang_rewriter.sh",
+    }),
+    visibility = ["//visibility:public"],
+)
+
 exports_files(
     srcs = ["deps.bzl"],
     visibility = ["//tools/private/update_deps:__pkg__"],
@@ -520,3 +538,15 @@ bzl_library(
     name = "whl_target_platforms_bzl",
     srcs = ["whl_target_platforms.bzl"],
 )
+
+bzl_library(
+    name = "venv_entry_point_bzl",
+    srcs = ["venv_entry_point.bzl"],
+    visibility = ["//visibility:public"],
+)
+
+bzl_library(
+    name = "venv_rewrite_shebang_bzl",
+    srcs = ["venv_rewrite_shebang.bzl"],
+    visibility = ["//visibility:public"],
+)

python/private/pypi/generate_whl_library_build_bazel.bzl

@@ -23,6 +23,7 @@ _RENDER = {
     "data_exclude": render.list,
     "dependencies": render.list,
     "dependencies_by_platform": lambda x: render.dict(x, value_repr = render.list),
+    "entry_points": render.dict_dict,
     "extras": render.list,
     "group_deps": render.list,
     "include": str,

python/private/pypi/venv_entry_point.bzl

@@ -0,0 +1,50 @@
+"""Rule for generating venv entry point scripts."""
+
+load("//python/private:attributes.bzl", "WINDOWS_CONSTRAINTS_ATTRS")
+load("//python/private:common.bzl", "is_windows_platform")
+load("//python/private:rule_builders.bzl", "ruleb")
+
+def _venv_entry_point_impl(ctx):
+    is_windows = is_windows_platform(ctx)
+
+    out_name = ctx.label.name
+    python_exe = ""
+    if is_windows:
+        out_name += ".bat"
+        python_exe = "pythonw.exe" if ctx.attr.group == "gui_scripts" else "python.exe"
+
+    out = ctx.actions.declare_file(out_name)
+
+    ctx.actions.expand_template(
+        template = ctx.file._template,
+        output = out,
+        substitutions = {
+            "{ATTRIBUTE}": ctx.attr.attribute,
+            "{MODULE}": ctx.attr.module,
+            "{PYTHON_EXE}": python_exe,
+        },
+        is_executable = True,
+    )
+
+    return [DefaultInfo(
+        files = depset([out]),
+        executable = out,
+    )]
+
+_builder = ruleb.Rule(
+    implementation = _venv_entry_point_impl,
+    executable = True,
+)
+_builder.attrs.update({
+    "attribute": attr.string(mandatory = False, doc = "The attribute to call"),
+    "extras": attr.string(mandatory = False, doc = "The extras for the entry point"),
+    "group": attr.string(mandatory = False, doc = "The entry point group (e.g. console_scripts)"),
+    "module": attr.string(mandatory = True, doc = "The module to import"),
+    "_template": attr.label(
+        default = Label("//python/private/pypi:venv_entry_point_template"),
+        allow_single_file = True,
+    ),
+})
+_builder.attrs.update(WINDOWS_CONSTRAINTS_ATTRS)
+
+venv_entry_point = _builder.build()

python/private/pypi/venv_entry_point_template.bat

@@ -0,0 +1,8 @@
+@setlocal enabledelayedexpansion & "%~dp0{PYTHON_EXE}" -x "%~f0" %* & exit /b !ERRORLEVEL!
+# -*- coding: utf-8 -*-
+import re
+import sys
+from {MODULE} import {ATTRIBUTE}
+if __name__ == "__main__":
+    sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
+    sys.exit({ATTRIBUTE}())

python/private/pypi/venv_entry_point_template.sh

@@ -0,0 +1,10 @@
+#!/bin/sh
+'''exec' "$(dirname "$0")/python3" "$0" "$@"
+' '''
+# -*- coding: utf-8 -*-
+import re
+import sys
+from {MODULE} import {ATTRIBUTE}
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit({ATTRIBUTE}())

python/private/pypi/venv_rewrite_shebang.bzl

@@ -0,0 +1,82 @@
+"""Rule for rewriting portable shebangs."""
+
+load("//python/private:attributes.bzl", "WINDOWS_CONSTRAINTS_ATTRS")
+load("//python/private:common.bzl", "is_windows_platform", "runfiles_root_path")
+load("//python/private:py_info.bzl", "PyInfoBuilder", "VenvSymlinkEntry", "VenvSymlinkKind")
+load("//python/private:rule_builders.bzl", "ruleb")
+
+def _venv_rewrite_shebang_impl(ctx):
+    is_windows = is_windows_platform(ctx)
+
+    out_name = ctx.label.name
+    if is_windows:
+        out_name += ".bat"
+
+    out_file = ctx.actions.declare_file(out_name)
+    in_file = ctx.file.src
+
+    action_args = ctx.actions.args()
+    rewriter_file = ctx.files._venv_shebang_rewriter[0]
+    inputs = depset([in_file, rewriter_file])
+
+    if rewriter_file.path.endswith(".ps1"):
+        action_exe = "powershell.exe"
+        action_args.add_all([
+            "-ExecutionPolicy",
+            "Bypass",
+            "-NoProfile",
+            "-File",
+            rewriter_file,
+        ])
+    else:
+        action_exe = ctx.attr._venv_shebang_rewriter[DefaultInfo].files_to_run
+
+    action_args.add(in_file)
+    action_args.add(out_file)
+    action_args.add("windows" if is_windows else "unix")
+
+    ctx.actions.run(
+        inputs = inputs,
+        outputs = [out_file],
+        executable = action_exe,
+        arguments = [action_args],
+        mnemonic = "PyVenvRewriteBin",
+        progress_message = "Rewriting venv bin script %{input}",
+        toolchain = None,
+    )
+
+    symlink = VenvSymlinkEntry(
+        kind = VenvSymlinkKind.BIN,
+        link_to_path = runfiles_root_path(ctx, out_file.short_path),
+        link_to_file = out_file,
+        venv_path = out_name,
+        package = ctx.attr.package,
+        version = ctx.attr.version,
+        files = depset([out_file]),
+    )
+    builder = PyInfoBuilder.new()
+    builder.venv_symlinks.add([symlink])
+    py_info = builder.build()
+
+    return [
+        DefaultInfo(files = depset([out_file]), executable = out_file),
+        py_info,
+    ]
+
+_builder = ruleb.Rule(
+    implementation = _venv_rewrite_shebang_impl,
+    executable = True,
+)
+_builder.attrs.update({
+    "package": attr.string(),
+    "src": attr.label(mandatory = True, allow_single_file = True),
+    "version": attr.string(),
+    "_venv_shebang_rewriter": attr.label(
+        default = "//python/private/pypi:venv_shebang_rewriter",
+        allow_files = True,
+        cfg = "exec",
+    ),
+})
+_builder.attrs.update(WINDOWS_CONSTRAINTS_ATTRS)
+
+venv_rewrite_shebang = _builder.build()

python/private/pypi/venv_shebang_rewriter.ps1

@@ -0,0 +1,44 @@
+[CmdletBinding()]
+param(
+    [Parameter(Position=0, Mandatory=$true)]
+    [string]$InFile,
+
+    [Parameter(Position=1, Mandatory=$true)]
+    [string]$OutFile,
+
+    [Parameter(Position=2, Mandatory=$true)]
+    [string]$TargetOs
+)
+
+$ErrorActionPreference = "Stop"
+
+$firstLine = Get-Content -Path $InFile -TotalCount 1 -ErrorAction SilentlyContinue
+$content = Get-Content -Path $InFile | Select-Object -Skip 1
+
+$Utf8NoBom = New-Object System.Text.UTF8Encoding $False
+
+if ($TargetOs -eq "windows") {
+    if ($firstLine -match "^#!pythonw") {
+        $pythonExe = "pythonw.exe"
+    } else {
+        $pythonExe = "python.exe"
+    }
+    # A Batch-Python polyglot. Batch executes the first line and exits,
+    # while Python (via -x) ignores the first line and executes the rest.
+    $wrapper = "@setlocal enabledelayedexpansion & `"%~dp0$pythonExe`" -x `"%~f0`" %* & exit /b !ERRORLEVEL!"
+    [System.IO.File]::WriteAllText($OutFile, $wrapper + "`r`n", $Utf8NoBom)
+} else {
+    # A Shell-Python polyglot. The shell executes the triple-quoted 'exec'
+    # command, re-running the script with python3 from the scripts directory.
+    # Python ignores the triple-quoted string and continues.
+    $wrapper = @"
+#!/bin/sh
+'''exec' "`$(dirname "`$0")/python3" "`$0" "`$@"
+' '''
+"@
+    [System.IO.File]::WriteAllText($OutFile, $wrapper + "`n", $Utf8NoBom)
+}
+
+if ($null -ne $content) {
+    [System.IO.File]::AppendAllLines($OutFile, [string[]]$content, $Utf8NoBom)
+}

python/private/pypi/venv_shebang_rewriter.sh

@@ -0,0 +1,27 @@
+#!/bin/sh
+set -eu
+
+IN="$1"
+OUT="$2"
+TARGET_OS="$3"
+
+FIRST_LINE=$(head -n 1 "$IN")
+
+if [ "$TARGET_OS" = "windows" ]; then
+  case "$FIRST_LINE" in
+    "#!pythonw"*) PYTHON_EXE="pythonw.exe" ;;
+    *)            PYTHON_EXE="python.exe" ;;
+  esac
+  # A Batch-Python polyglot. Batch executes the first line and exits,
+  # while Python (via -x) ignores the first line and executes the rest.
+  printf "@setlocal enabledelayedexpansion & \"%%~dp0$PYTHON_EXE\" -x \"%%~f0\" %%* & exit /b !ERRORLEVEL!\r\n" > "$OUT"
+else
+  printf "#!/bin/sh\n" > "$OUT"
+  # A Shell-Python polyglot. The shell executes the triple-quoted 'exec'
+  # command, re-running the script with python3 from the scripts directory.
+  # Python ignores the triple-quoted string and continues.
+  printf "'''exec' \"\$(dirname \"\$0\")/python3\" \"\$0\" \"\$@\"\n' '''\n" >> "$OUT"
+fi
+
+tail -n +2 "$IN" >> "$OUT"
+chmod +x "$OUT"