perf: py_venv.bzl optimizations

Author: dzbarsky

.bcr/patches/remove_dev_deps.patch

@@ -1,5 +1,5 @@
---- a/MODULE.bazel	2025-11-22 20:51:10
-+++ b/MODULE.bazel	2025-11-22 20:51:10
+--- a/MODULE.bazel	2025-11-23 11:09:03
++++ b/MODULE.bazel	2025-11-23 11:09:03
 @@ -39,558 +39,4 @@
 
  # HACK: In prod the includer's patch inserts the use_repo for multitool. This
@@ -373,7 +373,7 @@
 -bazel_dep(name = "xz", version = "5.4.5.bcr.5")
 -bazel_dep(name = "zstd", version = "1.5.7")
 -bazel_dep(name = "bzip2", version = "1.0.8.bcr.2")
--bazel_dep(name = "rules_rs", version = "0.0.7")
+-bazel_dep(name = "rules_rs", version = "0.0.15")
 -
 -crate = use_extension(
 -    "@rules_rs//rs:extensions.bzl",

MODULE.bazel

@@ -408,7 +408,7 @@ register_toolchains("@rust_toolchains//:all")
 bazel_dep(name = "xz", version = "5.4.5.bcr.5")
 bazel_dep(name = "zstd", version = "1.5.7")
 bazel_dep(name = "bzip2", version = "1.0.8.bcr.2")
-bazel_dep(name = "rules_rs", version = "0.0.7")
+bazel_dep(name = "rules_rs", version = "0.0.15")
 
 crate = use_extension(
     "@rules_rs//rs:extensions.bzl",

bazel/include/cargo.MODULE.bazel

@@ -4,7 +4,7 @@
 bazel_dep(name = "xz", version = "5.4.5.bcr.5")
 bazel_dep(name = "zstd", version = "1.5.7")
 bazel_dep(name = "bzip2", version = "1.0.8.bcr.2")
-bazel_dep(name = "rules_rs", version = "0.0.7")
+bazel_dep(name = "rules_rs", version = "0.0.15")
 
 crate = use_extension(
     "@rules_rs//rs:extensions.bzl",

py/private/py_venv/py_venv.bzl

@@ -28,6 +28,21 @@ def _interpreter_flags(ctx):
 
     return args
 
+# Forked from bazel-lib to avoid keeping `ctx` alive to execution phase
+def _to_rlocation_path(workspace_name, file):
+    if file.short_path.startswith("../"):
+        return file.short_path[3:]
+    else:
+        return workspace_name + "/" + file.short_path
+
+
+def _interpreter_path(workspace_name, py_toolchain):
+    return (
+        _to_rlocation_path(workspace_name, py_toolchain.python)
+        if py_toolchain.runfiles_interpreter
+        else py_toolchain.python.path
+    )
+
 # FIXME: This is derived directly from the py_binary.bzl rule and should really
 # be a layer on top of it if we can figure out flowing the data around. This is
 # PoC quality.
@@ -96,57 +111,73 @@ def _py_venv_base_impl(ctx):
 
     venv_name = ".{}".format(ctx.attr.name)
     venv_dir = ctx.actions.declare_directory(venv_name)
+    workspace_name = ctx.workspace_name
+
+    args = ctx.actions.args()
+    args.add_all([venv_dir], expand_directories = False, format_each="--location=%s")
+    args.add(py_shim.bin.bin, format="--venv-shim=%s")
+
+    # Post-bzlmod we need to record the current repository in case the
+    # user tries to consume a `py_venv_binary` across repo boundaries
+    # which could cause repo mapping to become relevant.
+    args.add(
+        ctx.label.repo_name or workspace_name,
+        format = "--repo=%s",
+    )
+    args.add_all(
+        [py_toolchain],
+        map_each=lambda py_toolchain: _interpreter_path(workspace_name, py_toolchain),
+        format_each = "--python=%s",
+        allow_closure = True)
+    args.add(site_packages_pth_file, format = "--pth-file=%s")
+    args.add(env_file, format = "--env-file=%s")
+    args.add(ctx.bin_dir.path, format = "--bin-dir=%s")
+    args.add(ctx.attr.package_collisions, format = "--collision-strategy=%s")
+    args.add(venv_name, format = "--venv-name=%s")
+    args.add(ctx.attr.mode, format = "--mode=%s")
+    args.add(
+        "{}.{}".format(
+            py_toolchain.interpreter_version_info.major,
+            py_toolchain.interpreter_version_info.minor,
+        ),
+        format = "--version=%s",
+    )
+    args.add(
+        "true" if ctx.attr.include_system_site_packages else "false",
+        format = "--include-system-site-packages=%s",
+    )
+    args.add(
+        "true" if ctx.attr.include_system_site_packages else "false",
+        format = "--include-user-site-packages=%s",
+    )
+
+    if ctx.attr.debug:
+        args.add("--debug")
 
     ctx.actions.run(
         executable = venv_tool,
-        arguments = [
-            "--location=" + venv_dir.path,
-            "--venv-shim=" + py_shim.bin.bin.path,
-            # Post-bzlmod we need to record the current repository in case the
-            # user tries to consume a `py_venv_binary` across repo boundaries
-            # which could cause repo mapping to become relevant.
-            "--repo=" + (ctx.label.repo_name or ctx.workspace_name),
-            "--python=" + (to_rlocation_path(ctx, py_toolchain.python) if py_toolchain.runfiles_interpreter else py_toolchain.python.path),
-            "--pth-file=" + site_packages_pth_file.path,
-            "--env-file=" + env_file.path,
-            "--bin-dir=" + ctx.bin_dir.path,
-            "--collision-strategy=" + ctx.attr.package_collisions,
-            "--venv-name=" + venv_name,
-            "--mode=" + ctx.attr.mode,
-            "--version={}.{}".format(
-                py_toolchain.interpreter_version_info.major,
-                py_toolchain.interpreter_version_info.minor,
-            ),
-            "--include-system-site-packages=" + ({
-                True: "true",
-                False: "false",
-            }[ctx.attr.include_system_site_packages]),
-            "--include-user-site-packages=" + ({
-                True: "true",
-                False: "false",
-            }[ctx.attr.include_system_site_packages]),
-        ] + (["--debug"] if ctx.attr.debug else []),
+        arguments = [args],
         inputs = rfs.merge_all([
-            ctx.runfiles(files = [
-                site_packages_pth_file,
-                env_file,
-                venv_tool,
-            ] + py_toolchain.files.to_list()),
+            ctx.runfiles(files = [site_packages_pth_file, env_file, venv_tool]),
             py_shim.default_info.default_runfiles,
         ]).files,
         outputs = [
             venv_dir,
         ],
+        execution_requirements = {
+            "supports-path-mapping": "1",
+        },
         # TODO: Is this right? The venv toolchain isn't quite in the right
         # configuration (target not exec) so we have to use a different source
         # of the target, but it is (logically) the venv toolchain.
         toolchain = VENV_TOOLCHAIN,
     )
 
     return venv_dir, rfs.merge_all([
-        ctx.runfiles(files = [
-            venv_dir,
-        ] + py_toolchain.files.to_list()),
+        ctx.runfiles(
+            files = [venv_dir],
+            transitive_files = py_toolchain.files,
+        ),
         ctx.attr._runfiles_lib[DefaultInfo].default_runfiles,
     ])