Add module_ctx.is_dev_dependency

Allows module extensions to determine whether a given tag represents a dev dependency.

Fixes #17101
Work towards #17908

Closes #17909.

PiperOrigin-RevId: 520645663
Change-Id: I3e3136a09d01d25fc706bcd0dfd7e53b6e7d5285

Author: fmeum

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java

@@ -24,6 +24,8 @@
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
 import java.util.Map;
 import javax.annotation.Nullable;
+import net.starlark.java.annot.Param;
+import net.starlark.java.annot.ParamType;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.eval.EvalException;
@@ -99,4 +101,20 @@ protected ImmutableMap<String, String> getRemoteExecProperties() throws EvalExce
   public StarlarkList<StarlarkBazelModule> getModules() {
     return modules;
   }
+
+  @StarlarkMethod(
+      name = "is_dev_dependency",
+      doc =
+          "Returns whether the given tag was specified on the result of a <a "
+              + "href=\"globals.html#use_extension\">use_extension</a> call with "
+              + "<code>devDependency = True</code>.",
+      parameters = {
+        @Param(
+            name = "tag",
+            doc = "A tag obtained from <a href=\"bazel_module.html#tags\">bazel_module.tags</a>.",
+            allowedTypes = {@ParamType(type = TypeCheckedTag.class)})
+      })
+  public boolean isDevDependency(TypeCheckedTag tag) {
+    return tag.isDevDependency();
+  }
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java

@@ -25,6 +25,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.docgen.annot.DocumentMethods;
+import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileGlobals.ModuleExtensionUsageBuilder.ModuleExtensionProxy;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import java.util.ArrayList;
@@ -63,7 +64,7 @@ public class ModuleFileGlobals {
   private final boolean ignoreDevDeps;
   private final Module.Builder module;
   private final Map<String, ModuleKey> deps = new LinkedHashMap<>();
-  private final List<ModuleExtensionProxy> extensionProxies = new ArrayList<>();
+  private final List<ModuleExtensionUsageBuilder> extensionUsageBuilders = new ArrayList<>();
   private final Map<String, ModuleOverride> overrides = new HashMap<>();
   private final Map<String, RepoNameUsage> repoNameUsages = new HashMap<>();
 
@@ -375,38 +376,37 @@ public void registerToolchains(Sequence<?> toolchainLabels) throws EvalException
       },
       useStarlarkThread = true)
   public ModuleExtensionProxy useExtension(
-      String extensionBzlFile, String extensionName, boolean devDependency, StarlarkThread thread)
-      throws EvalException {
-    ModuleExtensionProxy newProxy =
-        new ModuleExtensionProxy(extensionBzlFile, extensionName, thread.getCallerLocation());
+      String extensionBzlFile, String extensionName, boolean devDependency, StarlarkThread thread) {
+    ModuleExtensionUsageBuilder newUsageBuilder =
+        new ModuleExtensionUsageBuilder(
+            extensionBzlFile, extensionName, thread.getCallerLocation());
 
     if (ignoreDevDeps && devDependency) {
       // This is a no-op proxy.
-      return newProxy;
+      return newUsageBuilder.getProxy(devDependency);
     }
 
-    // Find an existing proxy object corresponding to this extension.
-    for (ModuleExtensionProxy proxy : extensionProxies) {
-      if (proxy.extensionBzlFile.equals(extensionBzlFile)
-          && proxy.extensionName.equals(extensionName)) {
-        return proxy;
+    // Find an existing usage builder corresponding to this extension.
+    for (ModuleExtensionUsageBuilder usageBuilder : extensionUsageBuilders) {
+      if (usageBuilder.extensionBzlFile.equals(extensionBzlFile)
+          && usageBuilder.extensionName.equals(extensionName)) {
+        return usageBuilder.getProxy(devDependency);
       }
     }
 
     // If no such proxy exists, we can just use a new one.
-    extensionProxies.add(newProxy);
-    return newProxy;
+    extensionUsageBuilders.add(newUsageBuilder);
+    return newUsageBuilder.getProxy(devDependency);
   }
 
-  @StarlarkBuiltin(name = "module_extension_proxy", documented = false)
-  class ModuleExtensionProxy implements Structure {
+  class ModuleExtensionUsageBuilder {
     private final String extensionBzlFile;
     private final String extensionName;
     private final Location location;
     private final HashBiMap<String, String> imports;
     private final ImmutableList.Builder<Tag> tags;
 
-    ModuleExtensionProxy(String extensionBzlFile, String extensionName, Location location) {
+    ModuleExtensionUsageBuilder(String extensionBzlFile, String extensionName, Location location) {
       this.extensionBzlFile = extensionBzlFile;
       this.extensionName = extensionName;
       this.location = location;
@@ -424,50 +424,69 @@ ModuleExtensionUsage buildUsage() {
           .build();
     }
 
-    void addImport(String localRepoName, String exportedName, Location location)
-        throws EvalException {
-      RepositoryName.validateUserProvidedRepoName(localRepoName);
-      RepositoryName.validateUserProvidedRepoName(exportedName);
-      addRepoNameUsage(localRepoName, "by a use_repo() call", location);
-      if (imports.containsValue(exportedName)) {
-        String collisionRepoName = imports.inverse().get(exportedName);
-        throw Starlark.errorf(
-            "The repo exported as '%s' by module extension '%s' is already imported at %s",
-            exportedName, extensionName, repoNameUsages.get(collisionRepoName).getWhere());
-      }
-      imports.put(localRepoName, exportedName);
+    /**
+     * Creates a proxy with the specified dev_dependency bit that shares accumulated imports and
+     * tags with all other such proxies, thus preserving their order across dev/non-dev deps.
+     */
+    ModuleExtensionProxy getProxy(boolean devDependency) {
+      return new ModuleExtensionProxy(devDependency);
     }
 
-    @Nullable
-    @Override
-    public Object getValue(String tagName) throws EvalException {
-      return new StarlarkValue() {
-        @StarlarkMethod(
-            name = "call",
-            selfCall = true,
-            documented = false,
-            extraKeywords = @Param(name = "kwargs"),
-            useStarlarkThread = true)
-        public void call(Dict<String, Object> kwargs, StarlarkThread thread) {
-          tags.add(
-              Tag.builder()
-                  .setTagName(tagName)
-                  .setAttributeValues(kwargs)
-                  .setLocation(thread.getCallerLocation())
-                  .build());
+    @StarlarkBuiltin(name = "module_extension_proxy", documented = false)
+    class ModuleExtensionProxy implements Structure {
+
+      private final boolean devDependency;
+
+      private ModuleExtensionProxy(boolean devDependency) {
+        this.devDependency = devDependency;
+      }
+
+      void addImport(String localRepoName, String exportedName, Location location)
+          throws EvalException {
+        RepositoryName.validateUserProvidedRepoName(localRepoName);
+        RepositoryName.validateUserProvidedRepoName(exportedName);
+        addRepoNameUsage(localRepoName, "by a use_repo() call", location);
+        if (imports.containsValue(exportedName)) {
+          String collisionRepoName = imports.inverse().get(exportedName);
+          throw Starlark.errorf(
+              "The repo exported as '%s' by module extension '%s' is already imported at %s",
+              exportedName, extensionName, repoNameUsages.get(collisionRepoName).getWhere());
         }
-      };
-    }
+        imports.put(localRepoName, exportedName);
+      }
 
-    @Override
-    public ImmutableCollection<String> getFieldNames() {
-      return ImmutableList.of();
-    }
+      @Nullable
+      @Override
+      public Object getValue(String tagName) throws EvalException {
+        return new StarlarkValue() {
+          @StarlarkMethod(
+              name = "call",
+              selfCall = true,
+              documented = false,
+              extraKeywords = @Param(name = "kwargs"),
+              useStarlarkThread = true)
+          public void call(Dict<String, Object> kwargs, StarlarkThread thread) {
+            tags.add(
+                Tag.builder()
+                    .setTagName(tagName)
+                    .setAttributeValues(kwargs)
+                    .setDevDependency(devDependency)
+                    .setLocation(thread.getCallerLocation())
+                    .build());
+          }
+        };
+      }
 
-    @Nullable
-    @Override
-    public String getErrorMessageForUnknownField(String field) {
-      return null;
+      @Override
+      public ImmutableCollection<String> getFieldNames() {
+        return ImmutableList.of();
+      }
+
+      @Nullable
+      @Override
+      public String getErrorMessageForUnknownField(String field) {
+        return null;
+      }
     }
   }
 
@@ -824,8 +843,8 @@ public Module buildModule() {
         .setDeps(ImmutableMap.copyOf(deps))
         .setOriginalDeps(ImmutableMap.copyOf(deps))
         .setExtensionUsages(
-            extensionProxies.stream()
-                .map(ModuleExtensionProxy::buildUsage)
+            extensionUsageBuilders.stream()
+                .map(ModuleExtensionUsageBuilder::buildUsage)
                 .collect(toImmutableList()))
         .build();
   }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Tag.java

@@ -33,6 +33,9 @@ public abstract class Tag {
   /** All keyword arguments supplied to the tag instance. */
   public abstract Dict<String, Object> getAttributeValues();
 
+  /** Whether this tag was created using a proxy created with dev_dependency = True. */
+  public abstract boolean isDevDependency();
+
   /** The source location in the module file where this tag was created. */
   public abstract Location getLocation();
 
@@ -48,6 +51,8 @@ public abstract static class Builder {
 
     public abstract Builder setAttributeValues(Dict<String, Object> value);
 
+    public abstract Builder setDevDependency(boolean value);
+
     public abstract Builder setLocation(Location value);
 
     public abstract Tag build();

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/TypeCheckedTag.java

@@ -34,10 +34,12 @@
 public class TypeCheckedTag implements Structure {
   private final TagClass tagClass;
   private final Object[] attrValues;
+  private final boolean devDependency;
 
-  private TypeCheckedTag(TagClass tagClass, Object[] attrValues) {
+  private TypeCheckedTag(TagClass tagClass, Object[] attrValues, boolean devDependency) {
     this.tagClass = tagClass;
     this.attrValues = attrValues;
+    this.devDependency = devDependency;
   }
 
   /** Creates a {@link TypeCheckedTag}. */
@@ -95,7 +97,15 @@ public static TypeCheckedTag create(TagClass tagClass, Tag tag, LabelConverter l
         attrValues[i] = Attribute.valueToStarlark(attr.getDefaultValueUnchecked());
       }
     }
-    return new TypeCheckedTag(tagClass, attrValues);
+    return new TypeCheckedTag(tagClass, attrValues, tag.isDevDependency());
+  }
+
+  /**
+   * Whether the tag was specified on an extension proxy created with <code>dev_dependency=True
+   * </code>.
+   */
+  public boolean isDevDependency() {
+    return devDependency;
   }
 
   @Override

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodTestUtil.java

@@ -283,6 +283,7 @@ public static TagClass createTagClass(Attribute... attrs) {
   public static class TestTagBuilder {
     private final Dict.Builder<String, Object> attrValuesBuilder = Dict.builder();
     private final String tagName;
+    private boolean devDependency = false;
 
     private TestTagBuilder(String tagName) {
       this.tagName = tagName;
@@ -294,11 +295,18 @@ public TestTagBuilder addAttr(String attrName, Object attrValue) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public TestTagBuilder setDevDependency() {
+      devDependency = true;
+      return this;
+    }
+
     public Tag build() {
       return Tag.builder()
           .setTagName(tagName)
           .setLocation(Location.BUILTIN)
           .setAttributeValues(attrValuesBuilder.buildImmutable())
+          .setDevDependency(devDependency)
           .build();
     }
   }

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionResolutionTest.java

@@ -446,7 +446,7 @@ public void multipleModules_devDependency() throws Exception {
         "  data_str = 'modules:'",
         "  for mod in ctx.modules:",
         "    for tag in mod.tags.tag:",
-        "      data_str += ' ' + tag.data",
+        "      data_str += ' ' + tag.data + ' ' + str(ctx.is_dev_dependency(tag))",
         "  data_repo(name='ext_repo',data=data_str)",
         "tag=tag_class(attrs={'data':attr.string()})",
         "ext=module_extension(implementation=_ext_impl,tag_classes={'tag':tag})");
@@ -457,7 +457,8 @@ public void multipleModules_devDependency() throws Exception {
     if (result.hasError()) {
       throw result.getError().getException();
     }
-    assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("modules: root bar@2.0");
+    assertThat(result.get(skyKey).getModule().getGlobal("data"))
+        .isEqualTo("modules: root True bar@2.0 False");
   }
 
   @Test
@@ -497,7 +498,7 @@ public void multipleModules_ignoreDevDependency() throws Exception {
         "  data_str = 'modules:'",
         "  for mod in ctx.modules:",
         "    for tag in mod.tags.tag:",
-        "      data_str += ' ' + tag.data",
+        "      data_str += ' ' + tag.data + ' ' + str(ctx.is_dev_dependency(tag))",
         "  data_repo(name='ext_repo',data=data_str)",
         "tag=tag_class(attrs={'data':attr.string()})",
         "ext=module_extension(implementation=_ext_impl,tag_classes={'tag':tag})");
@@ -511,7 +512,8 @@ public void multipleModules_ignoreDevDependency() throws Exception {
     if (result.hasError()) {
       throw result.getError().getException();
     }
-    assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("modules: bar@2.0");
+    assertThat(result.get(skyKey).getModule().getGlobal("data"))
+        .isEqualTo("modules: bar@2.0 False");
   }
 
   @Test