Allow native.package_relative_label in initializers

Work towards #21622 RELNOTES: `native.package_relative_label` can now be used in rule initializers. Closes #21740. PiperOrigin-RevId: 631546734 Change-Id: Id6eb237f0f79195b678bb954deaef071e1c45ab1
bazelbuild · May 7, 2024 · 8e7ef00 · 8e7ef00
1 parent 2330a11
commit 8e7ef00
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 3 deletions.
diff --git a/...main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java b/...main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
@@ -1272,6 +1272,9 @@ public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwarg
       BazelStarlarkContext bazelStarlarkContext = BazelStarlarkContext.fromOrFail(thread);
       try {
         thread.setThreadLocal(BazelStarlarkContext.class, null);
+        // Allow access to the LabelConverter to support native.package_relative_label() in an
+        // initializer.
+        thread.setThreadLocal(LabelConverter.class, pkgBuilder.getLabelConverter());
         thread.setUncheckedExceptionContext(() -> "an initializer");
 
         // We call all the initializers of the rule and its ancestor rules, proceeding from child to
@@ -1359,6 +1362,7 @@ public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwarg
           }
         }
       } finally {
+        thread.setThreadLocal(LabelConverter.class, null);
         bazelStarlarkContext.storeInThread(thread);
       }
 

diff --git a/src/main/java/com/google/devtools/build/lib/packages/StarlarkNativeModule.java b/src/main/java/com/google/devtools/build/lib/packages/StarlarkNativeModule.java
@@ -630,13 +630,18 @@ public String repoName(StarlarkThread thread) throws EvalException {
 
   @Override
   public Label packageRelativeLabel(Object input, StarlarkThread thread) throws EvalException {
-    BazelStarlarkContext.checkLoadingPhase(thread, "native.package_relative_label");
+    // In an initializer, BazelStarlarkContext isn't available, just the label converter.
+    LabelConverter labelConverter = thread.getThreadLocal(LabelConverter.class);
+    if (labelConverter == null) {
+      BazelStarlarkContext.checkLoadingPhase(thread, "native.package_relative_label");
+      Package.Builder pkgBuilder = Package.Builder.fromOrFail(thread, "package_relative_label()");
+      labelConverter = pkgBuilder.getLabelConverter();
+    }
     if (input instanceof Label inputLabel) {
       return inputLabel;
     }
     try {
-      Package.Builder pkgBuilder = Package.Builder.fromOrFail(thread, "package_relative_label()");
-      return pkgBuilder.getLabelConverter().convert((String) input);
+      return labelConverter.convert((String) input);
     } catch (LabelSyntaxException e) {
       throw Starlark.errorf("invalid label in native.package_relative_label: %s", e.getMessage());
     }

diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java
@@ -4320,6 +4320,44 @@ def impl(ctx):
     ev.assertContainsError("target 'my_target' not declared in package 'initializer_testing'");
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void initializer_nativeModule() throws Exception {
+    scratch.appendFile("MODULE.bazel", "module(name = 'my_mod', version = '1.2.3')");
+    scratch.file("initializer_testing/rules/BUILD");
+    scratch.file(
+        "initializer_testing/rules/b.bzl",
+        "MyInfo = provider()",
+        "def initializer(name, **kwargs):",
+        "  return {'props': {",
+        "    'package_relative_label': str(native.package_relative_label(':target')),",
+        "  }}",
+        "def impl(ctx): ",
+        "  return [MyInfo(props = ctx.attr.props)]",
+        "my_rule = rule(impl,",
+        "  initializer = initializer,",
+        "  attrs = {",
+        "    'props': attr.string_dict(),",
+        "  })");
+    scratch.file(
+        "initializer_testing/targets/BUILD",
+        "load('//initializer_testing/rules:b.bzl','my_rule')",
+        "my_rule(name = 'my_target')");
+
+    invalidatePackages();
+
+    ConfiguredTarget myTarget = getConfiguredTarget("//initializer_testing/targets:my_target");
+    StructImpl info =
+        (StructImpl)
+            myTarget.get(
+                new StarlarkProvider.Key(
+                    keyForBuild(Label.parseCanonical("//initializer_testing/rules:b.bzl")),
+                    "MyInfo"));
+
+    assertThat((Map<String, String>) info.getValue("props"))
+        .containsExactly("package_relative_label", "@@//initializer_testing/targets:target");
+  }
+
   private void scratchParentRule(String rule, String... ruleArgs) throws IOException {
     scratch.file("extend_rule_testing/parent/BUILD");
     scratch.file(