python: Add repo mapping to zip builds so runfiles lookup works correctly

To make runfiles lookups work correctly with bzlmod, a `_repo_mapping` file
is used to map apparent names to canonical names. This file is normally
automatically created by higher level code and injected into the runfiles, but
that happens after a rule has finished running. This means it isn't in the
runfiles object that py_binary itself sees when its constructing a zip file of
itself.

To fix, the zip file generates a repo mapping itself and puts it into the zip
file, thus allowing the runfiles library to find it when run from the zip file.
This requires a Java hook since the underlying `RepoMappingManifestAction`
isn't directly exposed to Starlark.

Fixes #17941

PiperOrigin-RevId: 532265478
Change-Id: I5f36fe58f043611481a7ed6ba19a7bbf5be4edf2

Author: rickeylev

src/main/java/com/google/devtools/build/lib/rules/python/BUILD

@@ -35,6 +35,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/patch_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/transition_factory",
         "//src/main/java/com/google/devtools/build/lib/analysis:file_provider",
+        "//src/main/java/com/google/devtools/build/lib/analysis:repo_mapping_manifest_action",
         "//src/main/java/com/google/devtools/build/lib/analysis:rule_definition_environment",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_provider",
@@ -44,6 +45,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/packages",
+        "//src/main/java/com/google/devtools/build/lib/packages/semantics",
         "//src/main/java/com/google/devtools/build/lib/rules/cpp",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant",

src/main/java/com/google/devtools/build/lib/rules/python/PyBuiltins.java

@@ -25,6 +25,7 @@
 import com.google.devtools.build.lib.analysis.AliasProvider;
 import com.google.devtools.build.lib.analysis.FileProvider;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
+import com.google.devtools.build.lib.analysis.RepoMappingManifestAction;
 import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesSupport;
 import com.google.devtools.build.lib.analysis.SingleRunfilesSupplier;
@@ -40,6 +41,7 @@
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.packages.StarlarkProvider;
+import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -66,6 +68,20 @@ protected PyBuiltins(Runfiles.EmptyFilesSupplier emptyFilesSupplier) {
     this.emptyFilesSupplier = emptyFilesSupplier;
   }
 
+  @StarlarkMethod(
+      name = "is_bzlmod_enabled",
+      doc = "Tells if bzlmod is enabled",
+      parameters = {
+        @Param(name = "ctx", positional = true, named = true, defaultValue = "unbound")
+      })
+  public boolean isBzlmodEnabled(StarlarkRuleContext starlarkCtx) {
+    return starlarkCtx
+        .getRuleContext()
+        .getAnalysisEnvironment()
+        .getStarlarkSemantics()
+        .getBool(BuildLanguageOptions.ENABLE_BZLMOD);
+  }
+
   @StarlarkMethod(
       name = "is_singleton_depset",
       doc = "Efficiently checks if the depset is a singleton.",
@@ -288,6 +304,28 @@ public Object newEmptyRunfilesWithMiddleman(
         .build();
   }
 
+  @StarlarkMethod(
+      name = "create_repo_mapping_manifest",
+      doc = "Write a repo_mapping file for the given runfiles",
+      parameters = {
+        @Param(name = "ctx", positional = false, named = true, defaultValue = "unbound"),
+        @Param(name = "runfiles", positional = false, named = true, defaultValue = "unbound"),
+        @Param(name = "output", positional = false, named = true, defaultValue = "unbound")
+      })
+  public void repoMappingAction(
+      StarlarkRuleContext starlarkCtx, Runfiles runfiles, Artifact repoMappingManifest) {
+    var ruleContext = starlarkCtx.getRuleContext();
+    ruleContext
+        .getAnalysisEnvironment()
+        .registerAction(
+            new RepoMappingManifestAction(
+                ruleContext.getActionOwner(),
+                repoMappingManifest,
+                ruleContext.getTransitivePackagesForRunfileRepoMappingManifest(),
+                runfiles.getAllArtifacts(),
+                ruleContext.getWorkspaceName()));
+  }
+
   @StarlarkMethod(
       name = "merge_runfiles_with_generated_inits_empty_files_supplier",
       doc =

src/main/starlark/builtins_bzl/common/python/py_executable_bazel.bzl

@@ -372,7 +372,24 @@ def _create_zip_file(ctx, *, output, original_nonzip_executable, executable_for_
             return None
 
     manifest.add_all(runfiles.files, map_each = map_zip_runfiles, allow_closure = True)
+
     inputs = [executable_for_zip_file]
+    if _py_builtins.is_bzlmod_enabled(ctx):
+        zip_repo_mapping_manifest = ctx.actions.declare_file(
+            output.basename + ".repo_mapping",
+            sibling = output,
+        )
+        _py_builtins.create_repo_mapping_manifest(
+            ctx = ctx,
+            runfiles = runfiles,
+            output = zip_repo_mapping_manifest,
+        )
+        manifest.add("{}/_repo_mapping={}".format(
+            _ZIP_RUNFILES_DIRECTORY_NAME,
+            zip_repo_mapping_manifest.path,
+        ))
+        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.

src/main/starlark/builtins_bzl/common/python/py_internal.bzl

@@ -28,6 +28,7 @@ py_internal = struct(
     add_py_extra_pseudo_action = _py_builtins.add_py_extra_pseudo_action,
     are_action_listeners_enabled = _py_builtins.are_action_listeners_enabled,
     copy_without_caching = _py_builtins.copy_without_caching,
+    create_repo_mapping_manifest = _py_builtins.create_repo_mapping_manifest,
     create_sources_only_manifest = _py_builtins.create_sources_only_manifest,
     declare_constant_metadata_file = _py_builtins.declare_constant_metadata_file,
     expand_location_and_make_variables = _py_builtins.expand_location_and_make_variables,
@@ -36,6 +37,7 @@ py_internal = struct(
     get_legacy_exernal_runfiles = _py_builtins.get_legacy_external_runfiles,
     get_rule_name = _py_builtins.get_rule_name,
     is_available_for = _is_available_for,
+    is_bzlmod_enabled = _py_builtins.is_bzlmod_enabled,
     is_singleton_depset = _py_builtins.is_singleton_depset,
     make_runfiles_respect_legacy_external_runfiles = _py_builtins.make_runfiles_respect_legacy_external_runfiles,
     merge_runfiles_with_generated_inits_empty_files_supplier = _py_builtins.merge_runfiles_with_generated_inits_empty_files_supplier,

src/test/shell/bazel/BUILD

@@ -917,6 +917,7 @@ sh_test(
         ":test-deps",
         "@bazel_tools//tools/bash/runfiles",
     ],
+    tags = ["requires-network"],
 )
 
 sh_test(

src/test/shell/bazel/python_version_test.sh

@@ -191,6 +191,37 @@ EOF
   expect_log "I am mockpy!"
 }
 
+# Verify that looking up runfiles that require repo mapping works
+function test_build_python_zip_bzlmod_repo_mapping_runfiles() {
+  cat > WORKSPACE
+  cat > MODULE.bazel << EOF
+module(name="pyzip")
+bazel_dep(name = "rules_python", version = "0.19.0")
+EOF
+  mkdir test
+  cat > test/BUILD << EOF
+py_binary(
+  name = "pybin",
+  srcs = ["pybin.py"],
+  deps = ["@rules_python//python/runfiles"],
+  data = ["data.txt"],
+)
+EOF
+  echo "data" > test/data.txt
+  cat > test/pybin.py << EOF
+from python.runfiles import runfiles
+rf = runfiles.Create()
+path = rf.Rlocation("pyzip/test/data.txt")
+with open(path, "r") as fp:
+  fp.read()
+EOF
+
+  bazel run --enable_bzlmod --build_python_zip //test:pybin &> $TEST_log || fail "bazel run failed"
+
+  unzip -p bazel-bin/test/pybin.zip runfiles/_repo_mapping > actual_repo_mapping
+  assert_contains ",pyzip,_main" actual_repo_mapping
+}
+
 # Test that running a zip app without RUN_UNDER_RUNFILES=1 removes the
 # temporary directory it creates
 function test_build_python_zip_cleans_up_temporary_module_space() {