Fix dict types in initializers

PiperOrigin-RevId: 597874318 Change-Id: I7d7135be863c2123242951635af5c10d370c7ef2
bazelbuild · Jan 12, 2024 · 48b975a · 48b975a
1 parent ad3dd45
commit 48b975a
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 1 deletion.
diff --git a/src/main/java/com/google/devtools/build/lib/packages/Type.java b/src/main/java/com/google/devtools/build/lib/packages/Type.java
@@ -606,9 +606,24 @@ public Map<KeyT, ValueT> convert(Object x, Object what, LabelConverter labelConv
     @Override
     public Object copyAndLiftStarlarkValue(
         Object x, Object what, @Nullable LabelConverter labelConverter) throws ConversionException {
-      return Dict.immutableCopyOf(convert(x, what, labelConverter));
+      if (!(x instanceof Map)) {
+        throw new ConversionException(this, x, what);
+      }
+      Map<?, ?> o = (Map<?, ?>) x;
+      // It's possible that #convert() calls transform non-equal keys into equal ones so we can't
+      // just use ImmutableMap.Builder() here (that throws on collisions).
+      LinkedHashMap<Object, Object> result = new LinkedHashMap<>();
+      for (Map.Entry<?, ?> elem : o.entrySet()) {
+        result.put(
+            keyType.copyAndLiftStarlarkValue(elem.getKey(), "dict key element", labelConverter),
+            valueType.copyAndLiftStarlarkValue(
+                elem.getValue(), "dict value element", labelConverter));
+      }
+      return Dict.immutableCopyOf(result);
     }
 
+
+
     @Override
     public Map<KeyT, ValueT> concat(Iterable<Map<KeyT, ValueT>> iterable) {
       Dict.Builder<KeyT, ValueT> output = new Dict.Builder<>();

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
@@ -86,6 +86,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
@@ -2993,6 +2994,73 @@ public void initializer_basic() throws Exception {
     assertThat((List<String>) info.getValue("deps")).containsExactly("@@//:initial", "@@//:added");
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void initializer_stringListDict() throws Exception {
+    scratch.file(
+        "initializer_testing/b.bzl",
+        "def initializer(**kwargs):",
+        "  return {}",
+        "MyInfo = provider()",
+        "def impl(ctx): ",
+        "  return [MyInfo(dict = ctx.attr.dict)]",
+        "my_rule = rule(impl,",
+        "  initializer = initializer,",
+        "  attrs = {",
+        "    'dict': attr.string_list_dict(),",
+        "  })");
+    scratch.file(
+        "initializer_testing/BUILD", //
+        "load(':b.bzl','my_rule')",
+        "my_rule(name = 'my_target', dict = {'k': ['val']})");
+
+    ConfiguredTarget myTarget = getConfiguredTarget("//initializer_testing:my_target");
+    StructImpl info =
+        (StructImpl)
+            myTarget.get(
+                new StarlarkProvider.Key(
+                    Label.parseCanonical("//initializer_testing:b.bzl"), "MyInfo"));
+
+    assertThat(((Map<String, List<String>>) info.getValue("dict")).keySet()).containsExactly("k");
+    assertThat(((Map<String, List<String>>) info.getValue("dict")).get("k")).containsExactly("val");
+  }
+
+  @Test
+  @SuppressWarnings("unchecked")
+  public void initializer_labelKeyedStringDict() throws Exception {
+    scratch.file(
+        "BUILD", //
+        "filegroup(name = 'key')");
+    scratch.file(
+        "initializer_testing/b.bzl",
+        "def initializer(**kwargs):",
+        "  return {}",
+        "MyInfo = provider()",
+        "def impl(ctx): ",
+        "  return [MyInfo(dict = ctx.attr.dict)]",
+        "my_rule = rule(impl,",
+        "  initializer = initializer,",
+        "  attrs = {",
+        "    'dict': attr.label_keyed_string_dict(),",
+        "  })");
+    scratch.file(
+        "initializer_testing/BUILD", //
+        "load(':b.bzl','my_rule')",
+        "my_rule(name = 'my_target', dict = {'//:key': 'val'})");
+
+    ConfiguredTarget myTarget = getConfiguredTarget("//initializer_testing:my_target");
+    ConfiguredTarget key = getConfiguredTarget("//:key");
+    StructImpl info =
+        (StructImpl)
+            myTarget.get(
+                new StarlarkProvider.Key(
+                    Label.parseCanonical("//initializer_testing:b.bzl"), "MyInfo"));
+
+    assertThat(((Map<ConfiguredTarget, String>) info.getValue("dict")).keySet())
+        .containsExactly(key);
+    assertThat(((Map<ConfiguredTarget, String>) info.getValue("dict")).get(key)).isEqualTo("val");
+  }
+
   @Test
   public void initializer_legacyAnyType() throws Exception {
     scratch.file(