Add set_select to set attr value to a select (bazelbuild#1153)

This code change adds a new command `set_select` which allows
setting the value of an attribute to a call to `select` and builds
the dict using the KVs passed to it.

Example call:
```
set_select <attr> <key_1> <value_1> <key_n> <value_n>
```

buildozer/README.md

@@ -119,6 +119,7 @@ Buildozer supports the following commands(`'command args'`):
   * `set_if_absent <attr> <value(s)>`: Sets the value of an attribute. If the
     attribute was already present, no action is taken.
   * `set kind <value>`: Set the target type to value.
+  * `set_select <attr> <key_1> <value_1> <key_n> <value_n>`
   * `copy <attr> <from_rule>`: Copies the value of `attr` between rules. If it
     exists in the `to_rule`, it will be overwritten.
   * `copy_no_overwrite <attr> <from_rule>`:  Copies the value of `attr` between

edit/buildozer.go

@@ -629,6 +629,38 @@ func cmdDictAdd(opts *Options, env CmdEnvironment) (*build.File, error) {
 	return env.File, nil
 }
 
+func cmdSetSelect(opts *Options, env CmdEnvironment) (*build.File, error) {
+	attr := env.Args[0]
+	args := env.Args[1:]
+
+	dict := &build.DictExpr{}
+
+	if len(args)%2 != 0 {
+		return nil, fmt.Errorf("no value passed for last key: %s", args[len(args)-1])
+	}
+	for i := 0; i < len(args); i += 2 {
+		key := args[i]
+		value := args[i+1]
+		var expr build.Expr
+		if IsList(attr) {
+			list := &build.ListExpr{}
+			if cur := DictionaryGet(dict, key); cur != nil {
+				list = cur.(*build.ListExpr)
+			}
+			AddValueToList(list, env.Pkg, getStringExpr(value, env.Pkg), !attributeMustNotBeSorted(env.Rule.Name(), attr))
+			expr = list
+		} else {
+			expr = getStringExpr(value, env.Pkg)
+		}
+		// Set overwrites previous values.
+		DictionarySet(dict, key, expr)
+	}
+	call := &build.CallExpr{List: []build.Expr{dict}}
+	call.X = &build.Ident{Name: "select"}
+	env.Rule.SetAttr(attr, call)
+	return env.File, nil
+}
+
 // cmdDictSet adds a key to a dict, overwriting any previous values.
 func cmdDictSet(opts *Options, env CmdEnvironment) (*build.File, error) {
 	attr := env.Args[0]
@@ -808,6 +840,7 @@ var AllCommands = map[string]CommandInfo{
 	"substitute":        {cmdSubstitute, true, 3, 3, "<attr> <old_regexp> <new_template>"},
 	"set":               {cmdSet, true, 1, -1, "<attr> <value(s)>"},
 	"set_if_absent":     {cmdSetIfAbsent, true, 1, -1, "<attr> <value(s)>"},
+	"set_select":        {cmdSetSelect, true, 1, -1, "<attr> <key_1> <value_1> <key_n> <value_n>"},
 	"copy":              {cmdCopy, true, 2, 2, "<attr> <from_rule>"},
 	"copy_no_overwrite": {cmdCopyNoOverwrite, true, 2, 2, "<attr> <from_rule>"},
 	"dict_add":          {cmdDictAdd, true, 2, -1, "<attr> <(key:value)(s)>"},

edit/buildozer_test.go

@@ -439,3 +439,94 @@ func TestCmdDictAddSet_missingColon(t *testing.T) {
 		})
 	}
 }
+
+func TestCmdSetSelect(t *testing.T) {
+	for i, tc := range []struct {
+		name      string
+		args      []string
+		buildFile string
+		expected  string
+	}{
+		{
+			name: "select_statment_doesn't_exist",
+			args: []string{
+				"args",                                   /* attr */
+				":use_ci_timeouts", "-test.timeout=123s", /* key, value */
+				":use_ci_timeouts", "-test.anotherFlag=flagValue", /* key, value */
+				"//conditions:default", "-test.timeout=789s", /* key, value */
+			},
+			buildFile: `foo(
+			name = "foo",
+)`,
+			expected: `foo(
+    name = "foo",
+    args = select({
+        ":use_ci_timeouts": [
+            "-test.timeout=123s",
+            "-test.anotherFlag=flagValue",
+        ],
+        "//conditions:default": ["-test.timeout=789s"],
+    }),
+)`},
+		{
+			name: "select_statment_exists",
+			args: []string{
+				"args",                                   /* attr */
+				":use_ci_timeouts", "-test.timeout=543s", /* key, value */
+				"//conditions:default", "-test.timeout=876s", /* key, value */
+			},
+			buildFile: `foo(
+    name = "foo",
+    args = select({
+        ":use_ci_timeouts": [
+            "-test.timeout=123s",
+            "-test.anotherFlag=flagValue",
+        ],
+        "//conditions:default": ["-test.timeout=789s"],
+    }),
+)`,
+			expected: `foo(
+    name = "foo",
+    args = select({
+        ":use_ci_timeouts": ["-test.timeout=543s"],
+        "//conditions:default": ["-test.timeout=876s"],
+    }),
+)`},
+		{
+			name: "attr_exists_but_not_select",
+			args: []string{
+				"args",                                   /* attr */
+				":use_ci_timeouts", "-test.timeout=543s", /* key, value */
+				"//conditions:default", "-test.timeout=876s", /* key, value */
+			},
+			buildFile: `foo(
+    name = "foo",
+    args = ["-test.timeout=123s"],
+)`,
+			expected: `foo(
+    name = "foo",
+    args = select({
+        ":use_ci_timeouts": ["-test.timeout=543s"],
+        "//conditions:default": ["-test.timeout=876s"],
+    }),
+)`},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			bld, err := build.Parse("BUILD", []byte(tc.buildFile))
+			if err != nil {
+				t.Error(err)
+			}
+			rl := bld.Rules("foo")[0]
+			env := CmdEnvironment{
+				File: bld,
+				Rule: rl,
+				Args: tc.args,
+			}
+			bld, _ = cmdSetSelect(NewOpts(), env)
+			got := strings.TrimSpace(string(build.Format(bld)))
+			if got != tc.expected {
+				t.Errorf("cmdSetSelect(%d):\ngot:\n%s\nexpected:\n%s", i, got, tc.expected)
+			}
+		})
+	}
+}