feat(zipapp): add content hash support to __main__-based invocations

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,20 +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 = {
-            "%EXTRACT_DIR%": paths.join(
-                (ctx.label.repo_name or "_main"),
-                ctx.label.package,
-                ctx.label.name,
-            ),
-            "%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
 
@@ -60,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.
@@ -75,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,
@@ -87,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")
@@ -106,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)
@@ -124,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}",
@@ -315,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

@@ -39,6 +39,7 @@
 _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")
 
@@ -197,7 +198,7 @@ def extract_zip(zip_path, dest_dir):
 # Create the runfiles tree by extracting the zip file
 def create_runfiles_root():
     if EXTRACT_ROOT:
-        extract_root = join(EXTRACT_ROOT, EXTRACT_DIR)
+        extract_root = join(EXTRACT_ROOT, EXTRACT_DIR, APP_HASH)
     else:
         extract_root = tempfile.mkdtemp("", "Bazel.runfiles_")
     extract_zip(dirname(__file__), extract_root)

tests/py_zipapp/system_python_zipapp_external_bootstrap_test.sh

@@ -29,5 +29,22 @@ if [[ ! -d "$RULES_PYTHON_EXTRACT_ROOT" ]]; then
   exit 1
 fi
 
+# The extract dir is _main/tests/py_zipapp/system_python_zipapp
+# The new structure should be $RULES_PYTHON_EXTRACT_ROOT/_main/tests/py_zipapp/system_python_zipapp/<hash>/runfiles
+# We check that there is a subdirectory under the expected extract dir.
+EXTRACT_DIR="_main/tests/py_zipapp/system_python_zipapp"
+if [[ ! -d "$RULES_PYTHON_EXTRACT_ROOT/$EXTRACT_DIR" ]]; then
+  echo "Error: Extract directory $RULES_PYTHON_EXTRACT_ROOT/$EXTRACT_DIR was not created!"
+  exit 1
+fi
+
+# Check for the extra hash component.
+# We use glob expansion to check for the expected depth.
+# Note: [ -d ... ] expands globs, while [[ -d ... ]] does not.
+if [ ! -d "$RULES_PYTHON_EXTRACT_ROOT/$EXTRACT_DIR"/*/runfiles ]; then
+  echo "Error: Could not find 'runfiles' directory at expected depth $RULES_PYTHON_EXTRACT_ROOT/$EXTRACT_DIR/*/runfiles"
+  exit 1
+fi
+
 echo "Running zipapp with extract root set a second time..."
 "$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"],
+)

tests/tools/zipapp/zip_main_maker_test.py

@@ -0,0 +1,101 @@
+import hashlib
+import os
+import tempfile
+import unittest
+from unittest import mock
+
+from tools.private.zipapp import zip_main_maker
+
+
+class ZipMainMakerTest(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp_dir.cleanup)
+
+    def test_creates_zip_main(self):
+        template_path = os.path.join(self.temp_dir.name, "template.py")
+        with open(template_path, "w", encoding="utf-8") as f:
+            f.write("hash=%APP_HASH%\nfoo=%FOO%\n")
+
+        output_path = os.path.join(self.temp_dir.name, "output.py")
+
+        file1_path = os.path.join(self.temp_dir.name, "file1.txt")
+        with open(file1_path, "wb") as f:
+            f.write(b"content1")
+
+        file2_path = os.path.join(self.temp_dir.name, "file2.txt")
+        with open(file2_path, "wb") as f:
+            f.write(b"content2")
+
+        # Add a symlink to test symlink hashing
+        symlink_path = os.path.join(self.temp_dir.name, "symlink.txt")
+        os.symlink(file1_path, symlink_path)
+
+        manifest_path = os.path.join(self.temp_dir.name, "manifest.txt")
+        with open(manifest_path, "w", encoding="utf-8") as f:
+            f.write(f"rf-file|0|file1.txt|{file1_path}\n")
+            f.write(f"rf-file|0|file2.txt|{file2_path}\n")
+            f.write(f"rf-symlink|1|symlink.txt|{symlink_path}\n")
+            f.write(f"rf-empty|empty_file.txt\n")
+
+        argv = [
+            "zip_main_maker.py",
+            "--template",
+            template_path,
+            "--output",
+            output_path,
+            "--substitution",
+            "%FOO%=bar",
+            "--hash_files_manifest",
+            manifest_path,
+        ]
+
+        with mock.patch("sys.argv", argv):
+            zip_main_maker.main()
+
+        # Calculate expected hash
+        h = hashlib.sha256()
+        line1 = f"rf-file|0|file1.txt|{file1_path}"
+        line2 = f"rf-file|0|file2.txt|{file2_path}"
+        line3 = f"rf-symlink|1|symlink.txt|{symlink_path}"
+        line4 = f"rf-empty|empty_file.txt"
+
+        # Sort lines like the program does
+        lines = sorted([line1, line2, line3, line4])
+        for line in lines:
+            parts = line.split("|")
+            if len(parts) > 1:
+                _, rest = line.split("|", 1)
+                h.update(rest.encode("utf-8"))
+            else:
+                h.update(line.encode("utf-8"))
+
+            type_ = parts[0]
+            if type_ == "rf-empty":
+                continue
+            if len(parts) >= 4:
+                is_symlink_str = parts[1]
+                path = parts[-1]
+                if not path:
+                    continue
+                if is_symlink_str == "-1":
+                    is_symlink = not os.path.exists(path)
+                else:
+                    is_symlink = is_symlink_str == "1"
+
+                if is_symlink:
+                    h.update(os.readlink(path).encode("utf-8"))
+                else:
+                    with open(path, "rb") as f:
+                        h.update(f.read())
+
+        expected_hash = h.hexdigest()
+
+        with open(output_path, "r", encoding="utf-8") as f:
+            content = f.read()
+
+        self.assertEqual(content, f"hash={expected_hash}\nfoo=bar\n")
+
+
+if __name__ == "__main__":
+    unittest.main()

tools/BUILD.bazel

@@ -31,6 +31,7 @@ filegroup(
         "wheelmaker.py",
         "//tools/launcher:distribution",
         "//tools/precompiler:distribution",
+        "//tools/private:distribution",
         "//tools/publish:distribution",
     ],
     visibility = ["//:__pkg__"],

tools/private/BUILD.bazel

@@ -0,0 +1,10 @@
+package(
+    default_visibility = ["//:__subpackages__"],
+)
+
+filegroup(
+    name = "distribution",
+    srcs = glob(["**"]) + [
+        "//tools/private/zipapp:distribution",
+    ],
+)

tools/private/zipapp/BUILD.bazel

@@ -34,3 +34,23 @@ py_library(
     name = "exe_zip_maker_lib",
     srcs = ["exe_zip_maker.py"],
 )
+
+py_interpreter_program(
+    name = "zip_main_maker",
+    main = "zip_main_maker.py",
+    visibility = [
+        # Not actually public. Only public so rules_python-generated toolchains
+        # are able to reference it.
+        "//visibility:public",
+    ],
+)
+
+py_library(
+    name = "zip_main_maker_lib",
+    srcs = ["zip_main_maker.py"],
+)
+
+filegroup(
+    name = "distribution",
+    srcs = glob(["**"]),
+)