feat(zipapp): support EXTRACT_ROOT env var for __main__.py invocations

python/private/zipapp/py_zipapp_rule.bzl

@@ -39,6 +39,11 @@ def _create_zipapp_main_py(ctx, py_runtime, py_executable, stage2_bootstrap):
             "%python_binary_actual%": python_binary_actual_path,
             "%stage2_bootstrap%": runfiles_root_path(ctx, stage2_bootstrap.short_path),
             "%workspace_name%": ctx.workspace_name,
+            "%EXTRACT_DIR%": paths.join(
+                (ctx.label.repo_name or "_main"),
+                ctx.label.package,
+                ctx.label.name,
+            ),
         },
     )
     return zip_main_py

python/private/zipapp/zip_main_template.py

@@ -22,6 +22,7 @@
 del sys.path[0]
 
 import os
+from os.path import join
 import shutil
 import subprocess
 import tempfile
@@ -35,6 +36,10 @@
 # executable to use.
 _PYTHON_BINARY_ACTUAL = "%python_binary_actual%"
 _WORKSPACE_NAME = "%workspace_name%"
+# relative path under EXTRACT_ROOT to extract to.
+EXTRACT_DIR = "%EXTRACT_DIR%"
+
+EXTRACT_ROOT = os.environ.get("RULES_PYTHON_EXTRACT_ROOT")
 
 
 def print_verbose(*args, mapping=None, values=None):
@@ -182,11 +187,14 @@ def extract_zip(zip_path, dest_dir):
 
 # Create the runfiles tree by extracting the zip file
 def create_runfiles_root():
-    temp_dir = tempfile.mkdtemp("", "Bazel.runfiles_")
-    extract_zip(os.path.dirname(__file__), temp_dir)
+    if EXTRACT_ROOT:
+        extract_root = join(EXTRACT_ROOT, EXTRACT_DIR)
+    else:
+        extract_root = tempfile.mkdtemp("", "Bazel.runfiles_")
+    extract_zip(os.path.dirname(__file__), extract_root)
     # IMPORTANT: Later code does `rm -fr` on dirname(runfiles_root) -- it's
     # important that deletion code be in sync with this directory structure
-    return os.path.join(temp_dir, "runfiles")
+    return os.path.join(extract_root, "runfiles")
 
 
 def execute_file(
@@ -223,18 +231,24 @@ def execute_file(
     # - When running in a zip file, we need to clean up the
     #   workspace after the process finishes so control must return here.
     try:
-        subprocess_argv = [python_program, main_filename] + args
+        subprocess_argv = [python_program]
+        if not EXTRACT_ROOT:
+            subprocess_argv.append(f"-XRULES_PYTHON_ZIP_DIR={os.path.dirname(runfiles_root)}")
+        subprocess_argv.append(main_filename)
+        subprocess_argv += args
         print_verbose("subprocess argv:", values=subprocess_argv)
         print_verbose("subprocess env:", mapping=env)
         print_verbose("subprocess cwd:", workspace)
         ret_code = subprocess.call(subprocess_argv, env=env, cwd=workspace)
         sys.exit(ret_code)
     finally:
-        # NOTE: dirname() is called because create_runfiles_root() creates a
-        # sub-directory within a temporary directory, and we want to remove the
-        # whole temporary directory.
-        ##shutil.rmtree(os.path.dirname(runfiles_root), True)
-        pass
+        if not EXTRACT_ROOT:
+            # NOTE: dirname() is called because create_runfiles_root() creates a
+            # sub-directory within a temporary directory, and we want to remove the
+            # whole temporary directory.
+            extract_root = os.path.dirname(runfiles_root)
+            print_verbose("cleanup: rmtree: ", extract_root)
+            shutil.rmtree(extract_root, True)
 
 
 def main():

tests/py_zipapp/BUILD.bazel

@@ -98,6 +98,17 @@ sh_test(
     toolchains = ["//python:current_py_toolchain"],
 )
 
+sh_test(
+    name = "extract_root_test",
+    srcs = ["extract_root_test.sh"],
+    data = [
+        ":system_python_zipapp",
+    ],
+    env = {
+        "ZIPAPP": "$(rootpath :system_python_zipapp)",
+    },
+)
+
 py_library(
     name = "some_dep",
     srcs = ["some_dep.py"],

tests/py_zipapp/extract_root_test.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+# Verify that the RULES_PYTHON_EXTRACT_ROOT env variable is respected.
+
+set -euo pipefail
+
+export RULES_PYTHON_EXTRACT_ROOT="${TEST_TMPDIR:-/tmp}/extract_root_test"
+
+echo "Running zipapp the first time..."
+"$ZIPAPP"
+
+# Verify that the directory was created
+if [[ ! -d "$RULES_PYTHON_EXTRACT_ROOT" ]]; then
+  echo "Error: Extract root directory $RULES_PYTHON_EXTRACT_ROOT was not created!"
+  exit 1
+fi
+
+# Run a second time to ensure it can re-extract successfully.
+echo "Running zipapp the second time..."
+"$ZIPAPP"
+
+echo "Success!"