perf: improve analysis performance by 95% for py_binary and py_test rules

Author: tobyh-canva

python/private/BUILD.bazel

@@ -14,6 +14,7 @@
 
 load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
 load("@bazel_skylib//rules:common_settings.bzl", "bool_setting")
+load("@rules_cc//cc:defs.bzl", "cc_binary")
 load("//python:py_binary.bzl", "py_binary")
 load("//python:py_library.bzl", "py_library")
 load(":print_toolchain_checksums.bzl", "print_toolchains_checksums")
@@ -829,6 +830,19 @@ py_binary(
     ],
 )
 
+# Used for py_executable rule
+# C++ wrapper for zipper to process Python zip manifests
+cc_binary(
+    name = "py_executable_zip_gen",
+    srcs = ["py_executable_zip_gen.cc"],
+    copts = select({
+        "@rules_cc//cc/compiler:msvc-cl": ["/std:c++17"],
+        "//conditions:default": ["-std=c++17"],
+    }),
+    data = ["@bazel_tools//tools/zip:zipper"],
+    deps = ["@bazel_tools//tools/cpp/runfiles"],
+)
+
 py_binary(
     name = "py_wheel_dist",
     srcs = ["py_wheel_dist.py"],

python/private/py_executable.bzl

@@ -67,8 +67,6 @@ load(":transition_labels.bzl", "TRANSITION_LABELS")
 load(":venv_runfiles.bzl", "create_venv_app_files")
 
 _py_builtins = py_internal
-_EXTERNAL_PATH_PREFIX = "external"
-_ZIP_RUNFILES_DIRECTORY_NAME = "runfiles"
 
 # Non-Google-specific attributes for executables
 # These attributes are for rules that accept Python sources.
@@ -236,7 +234,7 @@ accepting arbitrary Python versions.
         "_zipper": lambda: attrb.Label(
             cfg = "exec",
             executable = True,
-            default = "@bazel_tools//tools/zip:zipper",
+            default = ":py_executable_zip_gen",
         ),
     },
 )
@@ -380,9 +378,8 @@ def _create_executable(
     _create_zip_file(
         ctx,
         output = zip_file,
-        original_nonzip_executable = executable,
         zip_main = zip_main,
-        runfiles = runfiles_details.default_runfiles.merge(extra_runfiles),
+        runfiles = runfiles_details.runfiles_without_exe.merge(extra_runfiles),
     )
 
     extra_files_to_build = []
@@ -803,35 +800,16 @@ def _create_windows_exe_launcher(
         use_default_shell_env = True,
     )
 
-def _create_zip_file(ctx, *, output, original_nonzip_executable, zip_main, runfiles):
+def _create_zip_file(ctx, *, output, 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)
 
-    manifest = ctx.actions.args()
-    manifest.use_param_file("@%s", use_always = True)
-    manifest.set_param_file_format("multiline")
-
-    manifest.add("__main__.py={}".format(zip_main.path))
-    manifest.add("__init__.py=")
-    manifest.add(
-        "{}=".format(
-            _get_zip_runfiles_path("__init__.py", workspace_name, legacy_external_runfiles),
-        ),
-    )
-    for path in runfiles.empty_filenames.to_list():
-        manifest.add("{}=".format(_get_zip_runfiles_path(path, workspace_name, legacy_external_runfiles)))
-
-    def map_zip_runfiles(file):
-        if file != original_nonzip_executable and file != output:
-            return "{}={}".format(
-                _get_zip_runfiles_path(file.short_path, workspace_name, legacy_external_runfiles),
-                file.path,
-            )
-        else:
-            return None
-
-    manifest.add_all(runfiles.files, map_each = map_zip_runfiles, allow_closure = True)
+    args = ctx.actions.args()
+    args.add("--output", output)
+    args.add("--workspace-name", ctx.workspace_name)
+    args.add("--main-file", zip_main)
+    if legacy_external_runfiles:
+        args.add("--legacy-external-runfiles")
 
     inputs = [zip_main]
     if _py_builtins.is_bzlmod_enabled(ctx):
@@ -844,43 +822,30 @@ def _create_zip_file(ctx, *, output, original_nonzip_executable, zip_main, runfi
             runfiles = runfiles,
             output = zip_repo_mapping_manifest,
         )
-        manifest.add("{}/_repo_mapping={}".format(
-            _ZIP_RUNFILES_DIRECTORY_NAME,
-            zip_repo_mapping_manifest.path,
-        ))
+        args.add("--repo-mapping-manifest", zip_repo_mapping_manifest)
         inputs.append(zip_repo_mapping_manifest)
 
-    for artifact in runfiles.files.to_list():
-        # Don't include the original executable because it isn't used by the
-        # zip file, so no need to build it for the action.
-        # Don't include the zipfile itself because it's an output.
-        if artifact != original_nonzip_executable and artifact != output:
-            inputs.append(artifact)
-
-    zip_cli_args = ctx.actions.args()
-    zip_cli_args.add("cC")
-    zip_cli_args.add(output)
+    manifest = ctx.actions.args()
+    manifest.use_param_file("%s", use_always = True)
+    manifest.set_param_file_format("multiline")
+    manifest.add_all(runfiles.empty_filenames, map_each = _get_zip_empty_path_arg)
+    manifest.add_all(runfiles.files, map_each = _get_zip_path_arg)
 
     ctx.actions.run(
         executable = ctx.executable._zipper,
-        arguments = [zip_cli_args, manifest],
-        inputs = depset(inputs),
+        arguments = [args, manifest],
+        inputs = depset(inputs, transitive = [runfiles.files]),
         outputs = [output],
         use_default_shell_env = True,
         mnemonic = "PythonZipper",
         progress_message = "Building Python zip: %{label}",
     )
 
-def _get_zip_runfiles_path(path, workspace_name, legacy_external_runfiles):
-    if legacy_external_runfiles and path.startswith(_EXTERNAL_PATH_PREFIX):
-        zip_runfiles_path = paths.relativize(path, _EXTERNAL_PATH_PREFIX)
-    else:
-        # NOTE: External runfiles (artifacts in other repos) will have a leading
-        # path component of "../" so that they refer outside the main workspace
-        # directory and into the runfiles root. By normalizing, we simplify e.g.
-        # "workspace/../foo/bar" to simply "foo/bar".
-        zip_runfiles_path = paths.normalize("{}/{}".format(workspace_name, path))
-    return "{}/{}".format(_ZIP_RUNFILES_DIRECTORY_NAME, zip_runfiles_path)
+def _get_zip_empty_path_arg(file):
+    return "{}=".format(file)
+
+def _get_zip_path_arg(file):
+    return "{}={}".format(file.short_path, file.path)
 
 def _create_executable_zip_file(
         ctx,