Merge branch 'main' into chore.rm.pip.repository.annotations

Author: rickeylev

python/private/zipapp/py_zipapp_rule.bzl

@@ -18,7 +18,7 @@ def _is_symlink(f):
     else:
         return "-1"
 
-def _create_zipapp_main_py(ctx, py_runtime, py_executable, stage2_bootstrap):
+def _create_zipapp_main_py(ctx, py_runtime, py_executable, stage2_bootstrap, runfiles):
     venv_python_exe = py_executable.venv_python_exe
     if venv_python_exe:
         venv_python_exe_path = runfiles_root_path(ctx, venv_python_exe.short_path)
@@ -31,15 +31,40 @@ def _create_zipapp_main_py(ctx, py_runtime, py_executable, stage2_bootstrap):
         python_binary_actual_path = py_runtime.interpreter_path
 
     zip_main_py = ctx.actions.declare_file(ctx.label.name + ".zip_main.py")
-    ctx.actions.expand_template(
-        template = py_runtime.zip_main_template,
-        output = zip_main_py,
-        substitutions = {
-            "%python_binary%": venv_python_exe_path,
-            "%python_binary_actual%": python_binary_actual_path,
-            "%stage2_bootstrap%": runfiles_root_path(ctx, stage2_bootstrap.short_path),
-            "%workspace_name%": ctx.workspace_name,
-        },
+
+    args = ctx.actions.args()
+    args.add(py_runtime.zip_main_template, format = "--template=%s")
+    args.add(zip_main_py, format = "--output=%s")
+
+    args.add(
+        "%EXTRACT_DIR%=" + paths.join(
+            (ctx.label.repo_name or "_main"),
+            ctx.label.package,
+            ctx.label.name,
+        ),
+        format = "--substitution=%s",
+    )
+    args.add("%python_binary%=" + venv_python_exe_path, format = "--substitution=%s")
+    args.add("%python_binary_actual%=" + python_binary_actual_path, format = "--substitution=%s")
+    args.add("%stage2_bootstrap%=" + runfiles_root_path(ctx, stage2_bootstrap.short_path), format = "--substitution=%s")
+    args.add("%workspace_name%=" + ctx.workspace_name, format = "--substitution=%s")
+
+    hash_files_manifest = ctx.actions.args()
+    hash_files_manifest.use_param_file("--hash_files_manifest=%s", use_always = True)
+    hash_files_manifest.set_param_file_format("multiline")
+
+    inputs = builders.DepsetBuilder()
+    inputs.add(py_runtime.zip_main_template)
+    _build_manifest(ctx, hash_files_manifest, runfiles, inputs)
+
+    actions_run(
+        ctx,
+        executable = ctx.attr._zip_main_maker,
+        arguments = [args, hash_files_manifest],
+        inputs = inputs.build(),
+        outputs = [zip_main_py],
+        mnemonic = "PyZipAppCreateMainPy",
+        progress_message = "Generating zipapp __main__.py: %{label}",
     )
     return zip_main_py
 
@@ -55,9 +80,7 @@ def _map_zip_symlinks(entry):
 def _map_zip_root_symlinks(entry):
     return "rf-root-symlink|" + _is_symlink(entry.target_file) + "|" + entry.path + "|" + entry.target_file.path
 
-def _build_manifest(ctx, manifest, runfiles, zip_main):
-    manifest.add("regular|0|__main__.py|{}".format(zip_main.path))
-
+def _build_manifest(ctx, manifest, runfiles, inputs):
     manifest.add_all(
         # NOTE: Accessing runfiles.empty_filenames materializes them. A lambda
         # is used to defer that.
@@ -70,7 +93,10 @@ def _build_manifest(ctx, manifest, runfiles, zip_main):
     manifest.add_all(runfiles.symlinks, map_each = _map_zip_symlinks)
     manifest.add_all(runfiles.root_symlinks, map_each = _map_zip_root_symlinks)
 
-    inputs = [zip_main]
+    inputs.add(runfiles.files)
+    inputs.add([entry.target_file for entry in runfiles.symlinks.to_list()])
+    inputs.add([entry.target_file for entry in runfiles.root_symlinks.to_list()])
+
     zip_repo_mapping_manifest = maybe_create_repo_mapping(
         ctx = ctx,
         runfiles = runfiles,
@@ -82,8 +108,7 @@ def _build_manifest(ctx, manifest, runfiles, zip_main):
             zip_repo_mapping_manifest.path,
             format = "rf-root-symlink|0|_repo_mapping|%s",
         )
-        inputs.append(zip_repo_mapping_manifest)
-    return inputs
+        inputs.add(zip_repo_mapping_manifest)
 
 def _create_zip(ctx, py_runtime, py_executable, stage2_bootstrap):
     output = ctx.actions.declare_file(ctx.label.name + ".zip")
@@ -101,8 +126,17 @@ def _create_zip(ctx, py_runtime, py_executable, stage2_bootstrap):
 
     runfiles = runfiles.build(ctx)
 
-    zip_main = _create_zipapp_main_py(ctx, py_runtime, py_executable, stage2_bootstrap)
-    inputs = _build_manifest(ctx, manifest, runfiles, zip_main)
+    zip_main = _create_zipapp_main_py(
+        ctx,
+        py_runtime,
+        py_executable,
+        stage2_bootstrap,
+        runfiles,
+    )
+    inputs = builders.DepsetBuilder()
+    manifest.add("regular|0|__main__.py|{}".format(zip_main.path))
+    inputs.add(zip_main)
+    _build_manifest(ctx, manifest, runfiles, inputs)
 
     zipper_args = ctx.actions.args()
     zipper_args.add(output)
@@ -119,7 +153,7 @@ def _create_zip(ctx, py_runtime, py_executable, stage2_bootstrap):
         ctx,
         executable = ctx.attr._zipper,
         arguments = [manifest, zipper_args],
-        inputs = depset(inputs, transitive = [runfiles.files]),
+        inputs = inputs.build(),
         outputs = [output],
         mnemonic = "PyZipAppCreateZip",
         progress_message = "Reticulating zipapp archive: %{label} into %{output}",
@@ -310,6 +344,10 @@ Whether the output should be an executable zip file.
             "@platforms//os:windows",
         ],
     ),
+    "_zip_main_maker": attr.label(
+        cfg = "exec",
+        default = "//tools/private/zipapp:zip_main_maker",
+    ),
     "_zip_shell_template": attr.label(
         default = ":zip_shell_template",
         allow_single_file = True,

python/private/zipapp/zip_main_template.py

@@ -23,9 +23,11 @@
 
 import os
 import shutil
+import stat
 import subprocess
 import tempfile
 import zipfile
+from os.path import dirname, join, basename
 
 # runfiles-root-relative path
 _STAGE2_BOOTSTRAP = "%stage2_bootstrap%"
@@ -35,6 +37,12 @@
 # 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%"
+APP_HASH = "%APP_HASH%"
+
+EXTRACT_ROOT = os.environ.get("RULES_PYTHON_EXTRACT_ROOT")
+IS_WINDOWS = os.name == "nt"
 
 
 def print_verbose(*args, mapping=None, values=None):
@@ -61,10 +69,6 @@ def print_verbose(*args, mapping=None, values=None):
             print("bootstrap: stage 1:", *args, file=sys.stderr, flush=True)
 
 
-# Return True if running on Windows
-def is_windows():
-    return os.name == "nt"
-
 
 def get_windows_path_with_unc_prefix(path):
     """Adds UNC prefix after getting a normalized absolute Windows path.
@@ -75,7 +79,7 @@ def get_windows_path_with_unc_prefix(path):
 
     # No need to add prefix for non-Windows platforms.
     # And \\?\ doesn't work in python 2 or on mingw
-    if not is_windows() or sys.version_info[0] < 3:
+    if not IS_WINDOWS or sys.version_info[0] < 3:
         return path
 
     # Starting in Windows 10, version 1607(OS build 14393), MAX_PATH limitations have been
@@ -107,7 +111,7 @@ def has_windows_executable_extension(path):
 
 if (
     _PYTHON_BINARY_VENV
-    and is_windows()
+    and IS_WINDOWS
     and not has_windows_executable_extension(_PYTHON_BINARY_VENV)
 ):
     _PYTHON_BINARY_VENV = _PYTHON_BINARY_VENV + ".exe"
@@ -118,7 +122,7 @@ def search_path(name):
     search_path = os.getenv("PATH", os.defpath).split(os.pathsep)
     for directory in search_path:
         if directory:
-            path = os.path.join(directory, name)
+            path = join(directory, name)
             if os.path.isfile(path) and os.access(path, os.X_OK):
                 return path
     return None
@@ -139,7 +143,7 @@ def find_binary(runfiles_root, bin_name):
     # Use normpath() to convert slashes to os.sep on Windows.
     elif os.sep in os.path.normpath(bin_name):
         # Case 3: Path is relative to the repo root.
-        return os.path.join(runfiles_root, bin_name)
+        return join(runfiles_root, bin_name)
     else:
         # Case 4: Path has to be looked up in the search path.
         return search_path(bin_name)
@@ -161,10 +165,18 @@ def extract_zip(zip_path, dest_dir):
     dest_dir = get_windows_path_with_unc_prefix(dest_dir)
     with zipfile.ZipFile(zip_path) as zf:
         for info in zf.infolist():
+            file_path = os.path.abspath(join(dest_dir, info.filename))
+            # If the file exists, it might be a symlink or read-only file from a previous extraction.
+            # Unlink it first so zipfile.extract doesn't corrupt the symlink target or fail on read-only files.
+            if os.path.lexists(file_path) and not os.path.isdir(file_path):
+                try:
+                    os.unlink(file_path)
+                except OSError:
+                    # On Windows, unlinking a read-only file fails.
+                    os.chmod(file_path, stat.S_IWRITE)
+                    os.unlink(file_path)
+
             zf.extract(info, dest_dir)
-            # UNC-prefixed paths must be absolute/normalized. See
-            # https://docs.microsoft.com/en-us/windows/desktop/fileio/naming-a-file#maximum-path-length-limitation
-            file_path = os.path.abspath(os.path.join(dest_dir, info.filename))
             # The Unix st_mode bits (see "man 7 inode") are stored in the upper 16
             # bits of external_attr.
             attrs = info.external_attr >> 16
@@ -182,11 +194,21 @@ 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:
+        # Shorten the path for Windows in case long path support is disabled
+        if IS_WINDOWS:
+            hash_dir = APP_HASH[0:32]
+            extract_dir = basename(EXTRACT_DIR)
+            extract_root = join(EXTRACT_ROOT, extract_dir, hash_dir)
+        else:
+            extract_root = join(EXTRACT_ROOT, EXTRACT_DIR, APP_HASH)
+            extract_root = get_windows_path_with_unc_prefix(extract_root)
+    else:
+        extract_root = tempfile.mkdtemp("", "Bazel.runfiles_")
+    extract_zip(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 join(extract_root, "runfiles")
 
 
 def execute_file(
@@ -223,18 +245,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
-        print_verbose("subprocess argv:", values=subprocess_argv)
+        subprocess_argv = [python_program]
+        if not EXTRACT_ROOT:
+            subprocess_argv.append(f"-XRULES_PYTHON_ZIP_DIR={dirname(runfiles_root)}")
+        subprocess_argv.append(main_filename)
+        subprocess_argv += args
         print_verbose("subprocess env:", mapping=env)
         print_verbose("subprocess cwd:", workspace)
+        print_verbose("subprocess argv:", values=subprocess_argv)
         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 = dirname(runfiles_root)
+            print_verbose("cleanup: rmtree: ", extract_root)
+            shutil.rmtree(extract_root, True)
 
 
 def main():
@@ -254,7 +282,7 @@ def main():
 
     # The main Python source file.
     main_rel_path = _STAGE2_BOOTSTRAP
-    if is_windows():
+    if IS_WINDOWS:
         main_rel_path = main_rel_path.replace("/", os.sep)
 
     runfiles_root = create_runfiles_root()
@@ -266,7 +294,7 @@ def main():
     # See: https://docs.python.org/3.11/using/cmdline.html#envvar-PYTHONSAFEPATH
     new_env["PYTHONSAFEPATH"] = "1"
 
-    main_filename = os.path.join(runfiles_root, main_rel_path)
+    main_filename = join(runfiles_root, main_rel_path)
     main_filename = get_windows_path_with_unc_prefix(main_filename)
     assert os.path.exists(main_filename), (
         "Cannot exec() %r: file not found." % main_filename
@@ -276,7 +304,7 @@ def main():
     )
 
     if _PYTHON_BINARY_VENV:
-        python_program = os.path.join(runfiles_root, _PYTHON_BINARY_VENV)
+        python_program = join(runfiles_root, _PYTHON_BINARY_VENV)
         # When a venv is used, the `bin/python3` symlink may need to be created.
         # This case occurs when "create venv at runtime" or "resolve python at
         # runtime" modes are enabled.
@@ -288,7 +316,7 @@ def main():
                     "Program's venv binary not under runfiles: {python_program}"
                 )
             symlink_to = find_binary(runfiles_root, _PYTHON_BINARY_ACTUAL)
-            os.makedirs(os.path.dirname(python_program), exist_ok=True)
+            os.makedirs(dirname(python_program), exist_ok=True)
             try:
                 os.symlink(symlink_to, python_program)
             except OSError as e:
@@ -317,7 +345,7 @@ def main():
     # change directory to the right runfiles directory.
     # (So that the data files are accessible)
     if os.environ.get("RUN_UNDER_RUNFILES") == "1":
-        workspace = os.path.join(runfiles_root, _WORKSPACE_NAME)
+        workspace = join(runfiles_root, _WORKSPACE_NAME)
 
     sys.stdout.flush()
     execute_file(

tests/py_zipapp/system_python_zipapp_external_bootstrap_test.sh

@@ -13,6 +13,42 @@ fi
 ZIPAPP="${ZIPAPP/.exe/.zip}"
 
 export RULES_PYTHON_BOOTSTRAP_VERBOSE=1
+
 # We're testing the invocation of `__main__.py`, so we have to
 # manually pass the zipapp to python.
+echo "====================================================================="
+echo "Running zipapp using an automatic temp directory..."
+echo "====================================================================="
+"$PYTHON" "$ZIPAPP"
+
+echo
+echo
+
+echo "====================================================================="
+echo "Running zipapp with extract root set..."
+echo "====================================================================="
+export RULES_PYTHON_EXTRACT_ROOT="${TEST_TMPDIR:-/tmp}/extract_root_test"
+"$PYTHON" "$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
+
+# On windows, the path is shortened to just the basename to avoid long path errors.
+# Other platforms use the full path.
+# Note: [ -d ... ] expands globs, while [[ -d ... ]] does not.
+if [ -d "$RULES_PYTHON_EXTRACT_ROOT/_main/tests/py_zipapp/system_python_zipapp"/*/runfiles ]; then
+  echo "Found runfiles at $RULES_PYTHON_EXTRACT_ROOT/_main/tests/py_zipapp/system_python_zipapp/*/runfiles"
+elif [ -d "$RULES_PYTHON_EXTRACT_ROOT/system_python_zipapp"/*/runfiles ]; then
+  echo "Found runfiles at $RULES_PYTHON_EXTRACT_ROOT/system_python_zipapp/*/runfiles"
+else
+  echo "Error: Could not find 'runfiles' directory"
+  exit 1
+fi
+
+echo "====================================================================="
+echo "Running zipapp with extract root set a second time..."
+echo "====================================================================="
 "$PYTHON" "$ZIPAPP"

tests/tools/zipapp/BUILD.bazel

@@ -11,3 +11,9 @@ py_test(
     srcs = ["exe_zip_maker_test.py"],
     deps = ["//tools/private/zipapp:exe_zip_maker_lib"],
 )
+
+py_test(
+    name = "zip_main_maker_test",
+    srcs = ["zip_main_maker_test.py"],
+    deps = ["//tools/private/zipapp:zip_main_maker_lib"],
+)