From 6abbb6123ce57faacce627d8f9f94f01484e6967 Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Tue, 6 Dec 2022 13:01:29 +0100
Subject: [PATCH] Allow `map_each` to return `None` in
 `TemplateDict#add_joined`

Mimics the behavior of `Args#add_all`'s `map_each` callback and can be
used to skip over values.

Also improves the error message when the callback returns an unexpected
type by including the value.
---
 .../build/lib/analysis/starlark/TemplateDict.java      | 10 +++++-----
 .../build/lib/starlarkbuildapi/TemplateDictApi.java    |  6 +++---
 .../build/lib/starlark/StarlarkRuleContextTest.java    |  8 ++++----
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/TemplateDict.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/TemplateDict.java
index 23b2d9cf4b3012..153ef175f342c5 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/TemplateDict.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/TemplateDict.java
@@ -134,12 +134,12 @@ public String getValue() throws EvalException {
                     /*kwargs=*/ ImmutableMap.of());
             if (ret instanceof String) {
               parts.add((String) ret);
-              continue;
+            } 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));
             }
-            throw Starlark.errorf(
-                "Function provided to map_each must return a String, but returned type %s for key:"
-                    + " %s",
-                Starlark.type(ret), getKey());
           } catch (InterruptedException e) {
             // Report the error to the user, but the stack trace is not of use to them
             throw Starlark.errorf(
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/TemplateDictApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/TemplateDictApi.java
index 004af3001fee7a..e129f5e5fb740c 100644
--- a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/TemplateDictApi.java
+++ b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/TemplateDictApi.java
@@ -71,9 +71,9 @@ public interface TemplateDictApi extends StarlarkValue {
             named = true,
             positional = false,
             doc =
-                "A Starlark function accepting a single argument and returning a String. 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 or "
+                    + "<code>None</code>. This function is applied to each item of the depset "
+                    + "specified in the <code>values</code> parameter"),
         @Param(
             name = "uniquify",
             named = true,
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
index 64eb17fd1b8c2d..e5da90cc380b57 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
+++ b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
@@ -3626,7 +3626,7 @@ public void testTemplateExpansionComputedSubstitution() throws Exception {
     scratch.file(
         "test/rules.bzl",
         "def _artifact_to_basename(file):",
-        "  return file.basename",
+        "  return file.basename if file.basename != 'ignored.txt' else None",
         "",
         "def _undertest_impl(ctx):",
         "  template_dict = ctx.actions.template_dict()",
@@ -3657,7 +3657,7 @@ public void testTemplateExpansionComputedSubstitution() throws Exception {
         "undertest_rule(",
         "    name = 'undertest',",
         "    template = ':template.txt',",
-        "    srcs = ['foo.txt', 'bar.txt', 'baz.txt'],",
+        "    srcs = ['foo.txt', 'bar.txt', 'baz.txt', 'ignored.txt'],",
         ")",
         "testing_rule(",
         "    name = 'testing',",
@@ -3914,8 +3914,8 @@ public void testTemplateExpansionComputedSubstitutionMapEachBadReturnType() thro
     assertThat(evalException)
         .hasMessageThat()
         .isEqualTo(
-            "Function provided to map_each must return a String, but returned type Label for key:"
-                + " %files%");
+            "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>");
   }
 
   @Test