Allow TemplateDict#map_each callback to return a list of strings (bazelbuild#17306)

Returning a list of strings is needed to get correct `uniquify` semantics when a value may expand to multiple strings.

Work towards bazelbuild#1920

Closes bazelbuild#17008.

PiperOrigin-RevId: 495868960
Change-Id: I1f0ec560fd783846e916a07bb524dd2179f7c535

Co-authored-by: Fabian Meumertzheim <[email protected]>

Author: ShreeM01

src/main/java/com/google/devtools/build/lib/analysis/starlark/TemplateDict.java

@@ -28,6 +28,7 @@
 import java.util.List;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Mutability;
+import net.starlark.java.eval.Sequence;
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkCallable;
 import net.starlark.java.eval.StarlarkFunction;
@@ -121,11 +122,22 @@ public String getValue() throws EvalException {
                     /*kwargs=*/ ImmutableMap.of());
             if (ret instanceof String) {
               parts.add((String) ret);
+            } else if (ret instanceof Sequence) {
+              for (Object v : ((Sequence) ret)) {
+                if (!(v instanceof String)) {
+                  throw Starlark.errorf(
+                      "Function provided to map_each must return string, None, or list of strings,"
+                          + " but returned list containing element '%s' of type %s for key '%s' and"
+                          + " value: %s",
+                      v, Starlark.type(v), getKey(), val);
+                }
+                parts.add((String) v);
+              }
             } else if (ret != Starlark.NONE) {
               throw Starlark.errorf(
-                  "Function provided to map_each must return a String or None, but returned type "
-                      + "%s for key '%s' and value: %s",
-                  Starlark.type(ret), getKey(), Starlark.repr(val));
+                  "Function provided to map_each must return string, None, or list of strings, but "
+                      + "returned type %s for key '%s' and value: %s",
+                  Starlark.type(ret), getKey(), val);
             }
           } catch (InterruptedException e) {
             // Report the error to the user, but the stack trace is not of use to them

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/TemplateDictApi.java

@@ -70,9 +70,9 @@ public interface TemplateDictApi extends StarlarkValue {
             named = true,
             positional = false,
             doc =
-                "A Starlark function accepting a single argument and returning either a String or "
-                    + "<code>None</code>. This function is applied to each item of the depset "
-                    + "specified in the <code>values</code> parameter"),
+                "A Starlark function accepting a single argument and returning either a string, "
+                    + "<code>None</code>, or a list of strings. This function is applied to each "
+                    + "item of the depset specified in the <code>values</code> parameter"),
         @Param(
             name = "uniquify",
             named = true,

src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java

@@ -3740,6 +3740,68 @@ public void testTemplateExpansionComputedSubstitutionWithUniquify() throws Excep
     assertThat(substitutionsUnchecked).isEqualTo(ImmutableMap.of("exts", "txt%%log%%exe%%sh"));
   }
 
+  @Test
+  public void testTemplateExpansionComputedSubstitutionWithUniquifyAndListMapEach()
+      throws Exception {
+    setBuildLanguageOptions("--experimental_lazy_template_expansion");
+    scratch.file(
+        "test/rules.bzl",
+        "def _artifact_to_extension(file):",
+        "  if file.extension == 'sh':",
+        "    return [file.extension]",
+        "  return [file.extension, '.' + file.extension]",
+        "",
+        "def _undertest_impl(ctx):",
+        "  template_dict = ctx.actions.template_dict()",
+        "  template_dict.add_joined('exts', depset(ctx.files.srcs),",
+        "                           map_each = _artifact_to_extension,",
+        "                           uniquify = True,",
+        "                           join_with = '%%',",
+        "                          )",
+        "  ctx.actions.expand_template(output=ctx.outputs.out,",
+        "                              template=ctx.file.template,",
+        "                              computed_substitutions=template_dict,",
+        "                              )",
+        "undertest_rule = rule(",
+        "  implementation = _undertest_impl,",
+        "  outputs = {'out': '%{name}.txt'},",
+        "  attrs = {'template': attr.label(allow_single_file=True),",
+        "           'srcs':attr.label_list(allow_files=True)",
+        "           },",
+        "  _skylark_testable = True,",
+        ")",
+        testingRuleDefinition);
+    scratch.file("test/template.txt", "exts", "exts");
+    scratch.file(
+        "test/BUILD",
+        "load(':rules.bzl', 'undertest_rule', 'testing_rule')",
+        "undertest_rule(",
+        "    name = 'undertest',",
+        "    template = ':template.txt',",
+        "    srcs = ['foo.txt', 'bar.log', 'baz.txt', 'bak.exe', 'far.sh', 'boo.sh'],",
+        ")",
+        "testing_rule(",
+        "    name = 'testing',",
+        "    dep = ':undertest',",
+        ")");
+    StarlarkRuleContext ruleContext = createRuleContext("//test:testing");
+    setRuleContext(ruleContext);
+    ev.update("file", ev.eval("ruleContext.attr.dep.files.to_list()[0]"));
+    ev.update("action", ev.eval("ruleContext.attr.dep[Actions].by_file[file]"));
+
+    assertThat(ev.eval("type(action)")).isEqualTo("Action");
+
+    Object contentUnchecked = ev.eval("action.content");
+    assertThat(contentUnchecked).isInstanceOf(String.class);
+    assertThat(contentUnchecked)
+        .isEqualTo("txt%%.txt%%log%%.log%%exe%%.exe%%sh\ntxt%%.txt%%log%%.log%%exe%%.exe%%sh\n");
+
+    Object substitutionsUnchecked = ev.eval("action.substitutions");
+    assertThat(substitutionsUnchecked).isInstanceOf(Dict.class);
+    assertThat(substitutionsUnchecked)
+        .isEqualTo(ImmutableMap.of("exts", "txt%%.txt%%log%%.log%%exe%%.exe%%sh"));
+  }
+
   @Test
   public void testTemplateExpansionComputedSubstitutionDuplicateKeys() throws Exception {
     setBuildLanguageOptions("--experimental_lazy_template_expansion");
@@ -3912,8 +3974,60 @@ public void testTemplateExpansionComputedSubstitutionMapEachBadReturnType() thro
     assertThat(evalException)
         .hasMessageThat()
         .isEqualTo(
-            "Function provided to map_each must return a String or None, but returned "
-                + "type Label for key '%files%' and value: <source file test/template.txt>");
+            "Function provided to map_each must return string, None, or list of strings, "
+                + "but returned type Label for key '%files%' and value: "
+                + "File:[/workspace[source]]test/template.txt");
+  }
+
+  @Test
+  public void testTemplateExpansionComputedSubstitutionMapEachBadListReturnType() throws Exception {
+    setBuildLanguageOptions("--experimental_lazy_template_expansion");
+    scratch.file(
+        "test/rules.bzl",
+        "def file_to_owner_label(file):",
+        "  return [file.owner]",
+        "",
+        "def _undertest_impl(ctx):",
+        "  template_dict = ctx.actions.template_dict()",
+        "  template_dict.add_joined('%files%', depset(ctx.files.template),",
+        "                           map_each = file_to_owner_label,",
+        "                           join_with = '')",
+        "  ctx.actions.expand_template(output=ctx.outputs.out,",
+        "                              template=ctx.file.template,",
+        "                              computed_substitutions=template_dict,",
+        "                              )",
+        "undertest_rule = rule(",
+        "  implementation = _undertest_impl,",
+        "  outputs = {'out': '%{name}.txt'},",
+        "  attrs = {'template': attr.label(allow_single_file=True),},",
+        "  _skylark_testable = True,",
+        ")",
+        testingRuleDefinition);
+    scratch.file("test/template.txt");
+    scratch.file(
+        "test/BUILD",
+        "load(':rules.bzl', 'undertest_rule', 'testing_rule')",
+        "undertest_rule(",
+        "    name = 'undertest',",
+        "    template = ':template.txt',",
+        ")",
+        "testing_rule(",
+        "    name = 'testing',",
+        "    dep = ':undertest',",
+        ")");
+    StarlarkRuleContext ruleContext = createRuleContext("//test:testing");
+    setRuleContext(ruleContext);
+    ev.update("file", ev.eval("ruleContext.attr.dep.files.to_list()[0]"));
+    ev.update("action", ev.eval("ruleContext.attr.dep[Actions].by_file[file]"));
+
+    EvalException evalException =
+        assertThrows(EvalException.class, () -> ev.eval("action.content"));
+    assertThat(evalException)
+        .hasMessageThat()
+        .isEqualTo(
+            "Function provided to map_each must return string, None, or list of strings, "
+                + "but returned list containing element '//test:template.txt' of type Label for "
+                + "key '%files%' and value: File:[/workspace[source]]test/template.txt");
   }
 
   @Test