feat(coll): New coll.Set and coll.Unset functions (#2118)

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

5 files changed, 182 insertions, 0 deletions

diff --git a/docs-src/content/functions/coll.yml b/docs-src/content/functions/coll.yml
index 06027130..88ff1d98 100644
--- a/docs-src/content/functions/coll.yml
+++ b/docs-src/content/functions/coll.yml
@@ -461,3 +461,52 @@ funcs:
         {{ $keys := coll.Slice "foo" "baz" }}
         {{ coll.Omit $keys $data }}'
         map[bar:2]
+  - name: coll.Set
+    # released: v4.0.0
+    alias: set
+    description: |
+      Sets the given key to the given value in the given map.
+
+      The map is modified in place, and the modified map is returned.
+    pipeline: true
+    arguments:
+      - name: key
+        required: true
+        description: the key (string) to set
+      - name: value
+        required: true
+        description: the value to set
+      - name: map
+        required: true
+        description: the map to modify
+    examples:
+      - |
+        $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 }}
+        {{ coll.Set "baz" 3 $data }}'
+        map[bar:2 baz:3 foo:1]
+      - |
+        $ gomplate -i '{{ dict "foo" 1 | coll.Set "bar" 2 }}'
+        map[bar:2 foo:1]
+  - name: coll.Unset
+    # released: v4.0.0
+    alias: unset
+    description: |
+      Deletes the element with the specified key in the given map. If there is no such element, the map is returned unchanged.
+
+      The map is modified in place, and the modified map is returned.
+    pipeline: true
+    arguments:
+      - name: key
+        required: true
+        description: the key (string) to unset
+      - name: map
+        required: true
+        description: the map to modify
+    examples:
+      - |
+        $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+        {{ coll.Unset "bar" $data }}'
+        map[baz:3 foo:1]
+      - |
+        $ gomplate -i '{{ dict "foo" 1 "bar" 2 | coll.Unset "bar" }}'
+        map[foo:1]
diff --git a/docs/content/functions/coll.md b/docs/content/functions/coll.md
index 6448ac9b..a80d5a5f 100644
--- a/docs/content/functions/coll.md
+++ b/docs/content/functions/coll.md
@@ -706,3 +706,78 @@ $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
 {{ coll.Omit $keys $data }}'
 map[bar:2]
 ```
+
+## `coll.Set`_(unreleased)_
+**Unreleased:** _This function is in development, and not yet available in released builds of gomplate._
+
+**Alias:** `set`
+
+Sets the given key to the given value in the given map.
+
+The map is modified in place, and the modified map is returned.
+
+### Usage
+
+```
+coll.Set key value map
+```
+```
+map | coll.Set key value
+```
+
+### Arguments
+
+| name | description |
+|------|-------------|
+| `key` | _(required)_ the key (string) to set |
+| `value` | _(required)_ the value to set |
+| `map` | _(required)_ the map to modify |
+
+### Examples
+
+```console
+$ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 }}
+{{ coll.Set "baz" 3 $data }}'
+map[bar:2 baz:3 foo:1]
+```
+```console
+$ gomplate -i '{{ dict "foo" 1 | coll.Set "bar" 2 }}'
+map[bar:2 foo:1]
+```
+
+## `coll.Unset`_(unreleased)_
+**Unreleased:** _This function is in development, and not yet available in released builds of gomplate._
+
+**Alias:** `unset`
+
+Deletes the element with the specified key in the given map. If there is no such element, the map is returned unchanged.
+
+The map is modified in place, and the modified map is returned.
+
+### Usage
+
+```
+coll.Unset key map
+```
+```
+map | coll.Unset key
+```
+
+### Arguments
+
+| name | description |
+|------|-------------|
+| `key` | _(required)_ the key (string) to unset |
+| `map` | _(required)_ the map to modify |
+
+### Examples
+
+```console
+$ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+{{ coll.Unset "bar" $data }}'
+map[baz:3 foo:1]
+```
+```console
+$ gomplate -i '{{ dict "foo" 1 "bar" 2 | coll.Unset "bar" }}'
+map[foo:1]
+```
diff --git a/internal/funcs/coll.go b/internal/funcs/coll.go
index 945f77d3..0f87239d 100644
--- a/internal/funcs/coll.go
+++ b/internal/funcs/coll.go
@@ -33,6 +33,8 @@ func CreateCollFuncs(ctx context.Context) map[string]interface{} {
 	f["jsonpath"] = ns.JSONPath
 	f["jq"] = ns.JQ
 	f["flatten"] = ns.Flatten
+	f["set"] = ns.Set
+	f["unset"] = ns.Unset
 	return f
 }
 
@@ -221,3 +223,17 @@ func (CollFuncs) Omit(args ...interface{}) (map[string]interface{}, error) {
 	}
 	return coll.Omit(m, keys...), nil
 }
+
+// Set -
+func (CollFuncs) Set(key string, value interface{}, m map[string]interface{}) (map[string]interface{}, error) {
+	m[key] = value
+
+	return m, nil
+}
+
+// Unset -
+func (CollFuncs) Unset(key string, m map[string]interface{}) (map[string]interface{}, error) {
+	delete(m, key)
+
+	return m, nil
+}
diff --git a/internal/funcs/coll_test.go b/internal/funcs/coll_test.go
index fe0278d7..b075d078 100644
--- a/internal/funcs/coll_test.go
+++ b/internal/funcs/coll_test.go
@@ -216,3 +216,35 @@ func TestGoSlice(t *testing.T) {
 	assert.Equal(t, reflect.TypeOf([]string{}), out.Type())
 	assert.EqualValues(t, []string{"bar", "baz"}, out.Interface())
 }
+
+func TestCollFuncs_Set(t *testing.T) {
+	t.Parallel()
+
+	c := &CollFuncs{}
+
+	m := map[string]interface{}{"foo": "bar"}
+	out, err := c.Set("foo", "baz", m)
+	require.NoError(t, err)
+	assert.EqualValues(t, map[string]interface{}{"foo": "baz"}, out)
+
+	// m was modified so foo is now baz
+	out, err = c.Set("bar", "baz", m)
+	require.NoError(t, err)
+	assert.EqualValues(t, map[string]interface{}{"foo": "baz", "bar": "baz"}, out)
+}
+
+func TestCollFuncs_Unset(t *testing.T) {
+	t.Parallel()
+
+	c := &CollFuncs{}
+
+	m := map[string]interface{}{"foo": "bar"}
+	out, err := c.Unset("foo", m)
+	require.NoError(t, err)
+	assert.Empty(t, out)
+
+	// no-op
+	out, err = c.Unset("bar", m)
+	require.NoError(t, err)
+	assert.Empty(t, out)
+}
diff --git a/internal/tests/integration/collection_test.go b/internal/tests/integration/collection_test.go
index 45c33ad1..7ea15c1e 100644
--- a/internal/tests/integration/collection_test.go
+++ b/internal/tests/integration/collection_test.go
@@ -105,3 +105,13 @@ func TestColl_JQ(t *testing.T) {
 	inOutTest(t, `{{ coll.JQ ".foo" (dict "foo" 1 "bar" 2 "baz" 3) }}`, "1")
 	inOutTest(t, `{{ coll.Slice "one" 2 "three" 4.0 | jq ".[2]" }}`, `three`)
 }
+
+func TestColl_Set(t *testing.T) {
+	inOutTest(t, `{{ $dict := dict "foo" 1 }}{{ coll.Set "bar" 2 $dict }}`, "map[bar:2 foo:1]")
+	inOutTest(t, `{{ dict "foo" 1 | coll.Set "foo" 2 }}`, "map[foo:2]")
+}
+
+func TestColl_Unset(t *testing.T) {
+	inOutTest(t, `{{ $dict := dict "foo" 1 "bar" 2 }}{{ coll.Unset "bar" $dict }}`, "map[foo:1]")
+	inOutTest(t, `{{ dict "foo" 1 "bar" 2 | coll.Unset "foo" }}`, "map[bar:2]")
+}