Add dict_replace_if_equal command to buildozer. (#1274)

* Add dict_replace_if_equal command to buildozer. * Run gofmt -w and remove unnecessary equality check.
bazelbuild · Jun 26, 2024 · 9b954d0 · 9b954d0
1 parent 80f1f68
commit 9b954d0
Show file tree

Hide file tree

Showing 3 changed files with 178 additions and 29 deletions.
diff --git a/buildozer/README.md b/buildozer/README.md
@@ -147,6 +147,10 @@ Buildozer supports the following commands(`'command args'`):
   * `dict_set <attr> <(key:value)(s)>`:  Sets the value of a key for the dict
     attribute `attr`. If the key was already present, its old value is replaced.
   * `dict_remove <attr> <key(s)>`:  Deletes the key for the dict attribute `attr`.
+  * `dict_replace_if_equal <attr> <key> <old_value> <new_value>`: Replaces
+    `old_value` with `new_value` for key `key` in dictionary attribute `attr`.
+    If the key is not present in the dictionary, or does not have value
+    `old_value`, it will _not_ be updated.
   * `dict_list_add <attr> <key> <value(s)>`:  Adds value(s) to the list in the
     dict attribute `attr`.
   * `format`: Force formatting of all files, even if they were not changed by

diff --git a/edit/buildozer.go b/edit/buildozer.go
@@ -745,6 +745,35 @@ func cmdDictRemove(opts *Options, env CmdEnvironment) (*build.File, error) {
 	return env.File, nil
 }
 
+// cmdDictReplaceIfEqual updates a value in a dict if it is equal to a given value.
+func cmdDictReplaceIfEqual(opts *Options, env CmdEnvironment) (*build.File, error) {
+	attr := env.Args[0]
+	key := env.Args[1]
+	oldV := getStringValue(env.Args[2])
+	newV := getStringValue(env.Args[3])
+
+	thing := env.Rule.Attr(attr)
+	dictAttr, ok := thing.(*build.DictExpr)
+	if !ok {
+		return env.File, nil
+	}
+
+	prev := DictionaryGet(dictAttr, key)
+	if prev == nil {
+		return nil, fmt.Errorf("key '%s' not found in dict", key)
+	}
+	if e, ok := prev.(*build.StringExpr); ok {
+		if labels.Equal(e.Value, oldV, env.Pkg) {
+			DictionarySet(dictAttr, key, getStringExpr(newV, env.Pkg))
+		}
+	} else if e, ok := prev.(*build.Ident); ok {
+		if e.Name == oldV {
+			DictionarySet(dictAttr, key, getStringExpr(newV, env.Pkg))
+		}
+	}
+	return env.File, nil
+}
+
 // cmdDictListAdd adds an item to a list in a dict.
 func cmdDictListAdd(opts *Options, env CmdEnvironment) (*build.File, error) {
 	attr := env.Args[0]
@@ -880,35 +909,36 @@ type CommandInfo struct {
 // AllCommands associates the command names with their function and number
 // of arguments.
 var AllCommands = map[string]CommandInfo{
-	"add":               {cmdAdd, true, 2, -1, "<attr> <value(s)>"},
-	"new_load":          {cmdNewLoad, false, 1, -1, "<path> <[to=]from(s)>"},
-	"replace_load":      {cmdReplaceLoad, false, 1, -1, "<path> <[to=]symbol(s)>"},
-	"substitute_load":   {cmdSubstituteLoad, false, 2, 2, "<old_regexp> <new_template>"},
-	"comment":           {cmdComment, true, 1, 3, "<attr>? <value>? <comment>"},
-	"print_comment":     {cmdPrintComment, true, 0, 2, "<attr>? <value>?"},
-	"delete":            {cmdDelete, true, 0, 0, ""},
-	"fix":               {cmdFix, true, 0, -1, "<fix(es)>?"},
-	"move":              {cmdMove, true, 3, -1, "<old_attr> <new_attr> <value(s)>"},
-	"new":               {cmdNew, false, 2, 4, "<rule_kind> <rule_name> [(before|after) <relative_rule_name>]"},
-	"print":             {cmdPrint, true, 0, -1, "<attribute(s)>"},
-	"remove":            {cmdRemove, true, 1, -1, "<attr> <value(s)>"},
-	"remove_comment":    {cmdRemoveComment, true, 0, 2, "<attr>? <value>?"},
-	"remove_if_equal":   {cmdRemoveIfEqual, true, 2, 2, "<attr> <value>"},
-	"rename":            {cmdRename, true, 2, 2, "<old_attr> <new_attr>"},
-	"replace":           {cmdReplace, true, 3, 3, "<attr> <old_value> <new_value>"},
-	"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)>"},
-	"dict_set":          {cmdDictSet, true, 2, -1, "<attr> <(key:value)(s)>"},
-	"dict_remove":       {cmdDictRemove, true, 2, -1, "<attr> <key(s)>"},
-	"dict_list_add":     {cmdDictListAdd, true, 3, -1, "<attr> <key> <value(s)>"},
-	"use_repo_add":      {cmdUseRepoAdd, false, 2, -1, "([dev] <extension .bzl file> <extension name>|<use_extension variable name>) <repo(s)>"},
-	"use_repo_remove":   {cmdUseRepoRemove, false, 2, -1, "([dev] <extension .bzl file> <extension name>|<use_extension variable name>) <repo(s)>"},
-	"format":            {cmdFormat, false, 0, 0, ""},
+	"add":                   {cmdAdd, true, 2, -1, "<attr> <value(s)>"},
+	"new_load":              {cmdNewLoad, false, 1, -1, "<path> <[to=]from(s)>"},
+	"replace_load":          {cmdReplaceLoad, false, 1, -1, "<path> <[to=]symbol(s)>"},
+	"substitute_load":       {cmdSubstituteLoad, false, 2, 2, "<old_regexp> <new_template>"},
+	"comment":               {cmdComment, true, 1, 3, "<attr>? <value>? <comment>"},
+	"print_comment":         {cmdPrintComment, true, 0, 2, "<attr>? <value>?"},
+	"delete":                {cmdDelete, true, 0, 0, ""},
+	"fix":                   {cmdFix, true, 0, -1, "<fix(es)>?"},
+	"move":                  {cmdMove, true, 3, -1, "<old_attr> <new_attr> <value(s)>"},
+	"new":                   {cmdNew, false, 2, 4, "<rule_kind> <rule_name> [(before|after) <relative_rule_name>]"},
+	"print":                 {cmdPrint, true, 0, -1, "<attribute(s)>"},
+	"remove":                {cmdRemove, true, 1, -1, "<attr> <value(s)>"},
+	"remove_comment":        {cmdRemoveComment, true, 0, 2, "<attr>? <value>?"},
+	"remove_if_equal":       {cmdRemoveIfEqual, true, 2, 2, "<attr> <value>"},
+	"rename":                {cmdRename, true, 2, 2, "<old_attr> <new_attr>"},
+	"replace":               {cmdReplace, true, 3, 3, "<attr> <old_value> <new_value>"},
+	"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)>"},
+	"dict_set":              {cmdDictSet, true, 2, -1, "<attr> <(key:value)(s)>"},
+	"dict_remove":           {cmdDictRemove, true, 2, -1, "<attr> <key(s)>"},
+	"dict_replace_if_equal": {cmdDictReplaceIfEqual, true, 4, 4, "<attr> <key> <old_value> <new_value>"},
+	"dict_list_add":         {cmdDictListAdd, true, 3, -1, "<attr> <key> <value(s)>"},
+	"use_repo_add":          {cmdUseRepoAdd, false, 2, -1, "([dev] <extension .bzl file> <extension name>|<use_extension variable name>) <repo(s)>"},
+	"use_repo_remove":       {cmdUseRepoRemove, false, 2, -1, "([dev] <extension .bzl file> <extension name>|<use_extension variable name>) <repo(s)>"},
+	"format":                {cmdFormat, false, 0, 0, ""},
 }
 
 var readonlyCommands = map[string]bool{

diff --git a/edit/buildozer_test.go b/edit/buildozer_test.go
@@ -351,6 +351,121 @@ func TestCmdDictListAdd(t *testing.T) {
 	}
 }
 
+var dictReplaceIfEqualTests = []struct {
+	args      []string
+	buildFile string
+	expected  string
+}{
+	{[]string{
+		"attr", "key1", "value1", "value2",
+	},
+		`foo(
+		name = "foo",
+		attr = {"key1": "value1"},
+	)`,
+		`foo(
+    name = "foo",
+    attr = {"key1": "value2"},
+)`,
+	},
+	{[]string{
+		"attr", "key1", "value1", "value2",
+	},
+		`foo(
+		name = "foo",
+		attr = {"key1": "x"},
+	)`,
+		`foo(
+    name = "foo",
+    attr = {"key1": "x"},
+)`,
+	},
+	{[]string{
+		"attr", "key1", "value1", "value2",
+	},
+		`foo(
+		name = "foo",
+		attr = {
+      "key1": ["value1"],
+			"key2": ["value2"],
+	},
+	)`,
+		`foo(
+    name = "foo",
+    attr = {
+			"key1": ["value1"],
+			"key2": ["value2"],
+    },
+)`,
+	},
+	{[]string{
+		"attr", "key1", "value1", "value2",
+	},
+		`foo(
+		name = "foo",
+		attr = {
+			"key1": "value1",
+			"key2": "value2",
+	},
+	)`,
+		`foo(
+		name = "foo",
+		attr = {
+			"key1": "value2",
+			"key2": "value2",
+},
+)`,
+	},
+	{[]string{
+		"attr", "key1", "value1", "value2",
+	},
+		`foo(
+		name = "foo",
+		attr = {
+			"key1": "value1",
+			"key2": "x",
+	},
+	)`,
+		`foo(
+    name = "foo",
+    attr = {
+			"key1": "value2",
+			"key2": "x",
+		},
+)`,
+	},
+}
+
+func TestCmdDictReplaceIfEqual(t *testing.T) {
+	for i, tt := range dictReplaceIfEqualTests {
+		bld, err := build.Parse("BUILD", []byte(tt.buildFile))
+		if err != nil {
+			t.Error(err)
+			continue
+		}
+		expectedBld, err := build.Parse("BUILD", []byte(tt.expected))
+		if err != nil {
+			t.Error(err)
+			continue
+		}
+		rl := bld.Rules("foo")[0]
+		env := CmdEnvironment{
+			File: bld,
+			Rule: rl,
+			Args: tt.args,
+		}
+		bld, err = cmdDictReplaceIfEqual(NewOpts(), env)
+		if err != nil {
+			t.Errorf("cmdDictReplaceIfEqual(%d):\ngot error:\n%s", i, err)
+		}
+		got := strings.TrimSpace(string(build.Format(bld)))
+		expected := strings.TrimSpace(string(build.Format(expectedBld)))
+		if got != expected {
+			t.Errorf("cmdDictReplaceIfEqual(%d):\ngot:\n%s\nexpected:\n%s", i, got, tt.expected)
+		}
+	}
+}
+
 var substituteLoadsTests = []struct {
 	args      []string
 	buildFile string