Skip to content

Commit

Permalink
Add uniquify parameter to TemplateDict.add_joined
Browse files Browse the repository at this point in the history
The behavior is analogous to that of Args.add_joined.

Closes #16213.

PiperOrigin-RevId: 485571903
Change-Id: Id69de92d703d5202bfc7b50abfbbb4326441f242
  • Loading branch information
fmeum authored and copybara-github committed Nov 2, 2022
1 parent cdce638 commit 9d250ed
Show file tree
Hide file tree
Showing 3 changed files with 93 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.devtools.build.lib.analysis.actions.Substitution;
import com.google.devtools.build.lib.analysis.actions.Substitution.ComputedSubstitution;
import com.google.devtools.build.lib.collect.nestedset.Depset;
import com.google.devtools.build.lib.starlarkbuildapi.TemplateDictApi;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import net.starlark.java.eval.EvalException;
import net.starlark.java.eval.Mutability;
Expand Down Expand Up @@ -58,6 +60,7 @@ public TemplateDictApi addJoined(
Depset valuesSet,
String joinWith,
StarlarkCallable mapEach,
Boolean uniquify,
StarlarkThread thread)
throws EvalException {
if (mapEach instanceof StarlarkFunction) {
Expand All @@ -71,7 +74,7 @@ public TemplateDictApi addJoined(
}
}
substitutions.add(
new LazySubstitution(key, thread.getSemantics(), valuesSet, mapEach, joinWith));
new LazySubstitution(key, thread.getSemantics(), valuesSet, mapEach, uniquify, joinWith));
return this;
}

Expand All @@ -84,18 +87,21 @@ private static class LazySubstitution extends ComputedSubstitution {
private final StarlarkSemantics semantics;
private final Depset valuesSet;
private final StarlarkCallable mapEach;
private final boolean uniquify;
private final String joinWith;

public LazySubstitution(
String key,
StarlarkSemantics semantics,
Depset valuesSet,
StarlarkCallable mapEach,
boolean uniquify,
String joinWith) {
super(key);
this.semantics = semantics;
this.valuesSet = valuesSet;
this.mapEach = mapEach;
this.uniquify = uniquify;
this.joinWith = joinWith;
}

Expand Down Expand Up @@ -127,6 +133,19 @@ public String getValue() throws EvalException {
"Could not evaluate substitution for %s: %s", val, e.getMessage());
}
}
if (uniquify) {
// Stably deduplicate parts in-place.
int count = parts.size();
HashSet<String> seen = Sets.newHashSetWithExpectedSize(count);
int addIndex = 0;
for (int i = 0; i < count; ++i) {
String val = parts.get(i);
if (seen.add(val)) {
parts.set(addIndex++, val);
}
}
parts = parts.subList(0, addIndex);
}
return Joiner.on(joinWith).join(parts);
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,9 +73,24 @@ public interface TemplateDictApi extends StarlarkValue {
"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"),
@Param(
name = "uniquify",
named = true,
positional = false,
defaultValue = "False",
doc =
"If true, duplicate strings derived from <code>values</code> will be omitted. Only "
+ "the first occurrence of each string will remain. Usually this feature is "
+ "not needed because depsets already omit duplicates, but it can be useful "
+ "if <code>map_each</code> emits the same string for multiple items."),
},
useStarlarkThread = true)
TemplateDictApi addJoined(
String key, Depset values, String joinWith, StarlarkCallable mapEach, StarlarkThread thread)
String key,
Depset values,
String joinWith,
StarlarkCallable mapEach,
Boolean uniquify,
StarlarkThread thread)
throws EvalException;
}
Original file line number Diff line number Diff line change
Expand Up @@ -3683,6 +3683,63 @@ public void testTemplateExpansionComputedSubstitution() throws Exception {
"td_files_key", "foo.txt%%bar.txt%%baz.txt"));
}

@Test
public void testTemplateExpansionComputedSubstitutionWithUniquify() throws Exception {
setBuildLanguageOptions("--experimental_lazy_template_expansion");
scratch.file(
"test/rules.bzl",
"def _artifact_to_extension(file):",
" return 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%%log%%exe%%sh\ntxt%%log%%exe%%sh\n");

Object substitutionsUnchecked = ev.eval("action.substitutions");
assertThat(substitutionsUnchecked).isInstanceOf(Dict.class);
assertThat(substitutionsUnchecked).isEqualTo(ImmutableMap.of("exts", "txt%%log%%exe%%sh"));
}

@Test
public void testTemplateExpansionComputedSubstitutionDuplicateKeys() throws Exception {
setBuildLanguageOptions("--experimental_lazy_template_expansion");
Expand Down

0 comments on commit 9d250ed

Please sign in to comment.