New functions coll.Pick and coll.Omit

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

7 files changed, 336 insertions, 0 deletions

diff --git a/coll/coll.go b/coll/coll.go
index 284d5f29..921abd53 100644
--- a/coll/coll.go
+++ b/coll/coll.go
@@ -162,6 +162,40 @@ func Merge(dst map[string]interface{}, srcs ...map[string]interface{}) (map[stri
 	return dst, nil
 }
 
+// returns whether or not a contains v
+func contains(v string, a []string) bool {
+	for _, n := range a {
+		if n == v {
+			return true
+		}
+	}
+	return false
+}
+
+// Omit returns a new map without any entries that have the
+// given keys (inverse of Pick).
+func Omit(in map[string]interface{}, keys ...string) map[string]interface{} {
+	out := map[string]interface{}{}
+	for k, v := range in {
+		if !contains(k, keys) {
+			out[k] = v
+		}
+	}
+	return out
+}
+
+// Pick returns a new map with any entries that have the
+// given keys (inverse of Omit).
+func Pick(in map[string]interface{}, keys ...string) map[string]interface{} {
+	out := map[string]interface{}{}
+	for k, v := range in {
+		if contains(k, keys) {
+			out[k] = v
+		}
+	}
+	return out
+}
+
 func copyMap(m map[string]interface{}) map[string]interface{} {
 	n := map[string]interface{}{}
 	for k, v := range m {
diff --git a/coll/coll_test.go b/coll/coll_test.go
index 6110a78d..ee6f7a2d 100644
--- a/coll/coll_test.go
+++ b/coll/coll_test.go
@@ -565,3 +565,48 @@ func BenchmarkFlatten(b *testing.B) {
 		}
 	}
 }
+
+func TestOmit(t *testing.T) {
+	in := map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+		"":    "baz",
+	}
+	assert.EqualValues(t, in, Omit(in, "baz"))
+
+	expected := map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+	}
+	assert.EqualValues(t, expected, Omit(in, ""))
+
+	expected = map[string]interface{}{
+		"": "baz",
+	}
+	assert.EqualValues(t, expected, Omit(in, "foo", "bar"))
+
+	assert.EqualValues(t, map[string]interface{}{}, Omit(in, "foo", "bar", ""))
+}
+
+func TestPick(t *testing.T) {
+	in := map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+		"":    "baz",
+	}
+	expected := map[string]interface{}{}
+	assert.EqualValues(t, expected, Pick(in, "baz"))
+
+	expected = map[string]interface{}{
+		"": "baz",
+	}
+	assert.EqualValues(t, expected, Pick(in, ""))
+
+	expected = map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+	}
+	assert.EqualValues(t, expected, Pick(in, "foo", "bar"))
+
+	assert.EqualValues(t, in, Pick(in, "foo", "bar", ""))
+}
diff --git a/docs-src/content/functions/coll.yml b/docs-src/content/functions/coll.yml
index 8a9917f1..cf267fe1 100644
--- a/docs-src/content/functions/coll.yml
+++ b/docs-src/content/functions/coll.yml
@@ -301,3 +301,47 @@ funcs:
         {{ $src2 := dict "foo" 3 "bar" 5 }}
         {{ coll.Merge $dst $src1 $src2 }}'
         map[foo:1 bar:5 baz:4]
+  - name: coll.Pick
+    description: |
+      Given a map, returns a new map with any entries that have the given keys.
+
+      All keys are converted to strings.
+
+      This is the inverse of [`coll.Omit`](#coll-omit).
+
+      _Note that this function does not modify the input._
+    pipeline: true
+    arguments:
+      - name: keys...
+        required: true
+        description: the keys to match
+      - name: map
+        required: true
+        description: the map to pick from
+    examples:
+      - |
+        $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+        {{ coll.Pick "foo" "baz" $data }}'
+        map[baz:3 foo:1]
+  - name: coll.Omit
+    description: |
+      Given a map, returns a new map without any entries that have the given keys.
+
+      All keys are converted to strings.
+
+      This is the inverse of [`coll.Pic`](#coll-pick).
+
+      _Note that this function does not modify the input._
+    pipeline: true
+    arguments:
+      - name: keys...
+        required: true
+        description: the keys to match
+      - name: map
+        required: true
+        description: the map to omit from
+    examples:
+      - |
+        $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+        {{ coll.Omit "foo" "baz" $data }}'
+        map[bar:2]
diff --git a/docs/content/functions/coll.md b/docs/content/functions/coll.md
index 6ce9645b..31878333 100644
--- a/docs/content/functions/coll.md
+++ b/docs/content/functions/coll.md
@@ -484,3 +484,71 @@ $ gomplate -i '{{ $dst := dict "foo" 1 "bar" 2 }}
 {{ coll.Merge $dst $src1 $src2 }}'
 map[foo:1 bar:5 baz:4]
 ```
+
+## `coll.Pick`
+
+Given a map, returns a new map with any entries that have the given keys.
+
+All keys are converted to strings.
+
+This is the inverse of [`coll.Omit`](#coll-omit).
+
+_Note that this function does not modify the input._
+
+### Usage
+
+```go
+coll.Pick keys... map
+```
+```go
+map | coll.Pick keys...
+```
+
+### Arguments
+
+| name | description |
+|------|-------------|
+| `keys...` | _(required)_ the keys to match |
+| `map` | _(required)_ the map to pick from |
+
+### Examples
+
+```console
+$ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+{{ coll.Pick "foo" "baz" $data }}'
+map[baz:3 foo:1]
+```
+
+## `coll.Omit`
+
+Given a map, returns a new map without any entries that have the given keys.
+
+All keys are converted to strings.
+
+This is the inverse of [`coll.Pic`](#coll-pick).
+
+_Note that this function does not modify the input._
+
+### Usage
+
+```go
+coll.Omit keys... map
+```
+```go
+map | coll.Omit keys...
+```
+
+### Arguments
+
+| name | description |
+|------|-------------|
+| `keys...` | _(required)_ the keys to match |
+| `map` | _(required)_ the map to omit from |
+
+### Examples
+
+```console
+$ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+{{ coll.Omit "foo" "baz" $data }}'
+map[bar:2]
+```
diff --git a/funcs/coll.go b/funcs/coll.go
index 1e110b8b..570b6544 100644
--- a/funcs/coll.go
+++ b/funcs/coll.go
@@ -129,3 +129,42 @@ func (f *CollFuncs) Flatten(args ...interface{}) ([]interface{}, error) {
 	}
 	return coll.Flatten(list, depth)
 }
+
+func pickOmitArgs(args ...interface{}) (map[string]interface{}, []string, error) {
+	if len(args) <= 1 {
+		return nil, nil, errors.Errorf("wrong number of args: wanted 2 or more, got %d", len(args))
+	}
+
+	m, ok := args[len(args)-1].(map[string]interface{})
+	if !ok {
+		return nil, nil, errors.Errorf("wrong map type: must be map[string]interface{}, got %T", args[len(args)-1])
+	}
+
+	keys := make([]string, len(args)-1)
+	for i, v := range args[0 : len(args)-1] {
+		k, ok := v.(string)
+		if !ok {
+			return nil, nil, errors.Errorf("wrong key type: must be string, got %T (%+v)", args[i], args[i])
+		}
+		keys[i] = k
+	}
+	return m, keys, nil
+}
+
+// Pick -
+func (f *CollFuncs) Pick(args ...interface{}) (map[string]interface{}, error) {
+	m, keys, err := pickOmitArgs(args...)
+	if err != nil {
+		return nil, err
+	}
+	return coll.Pick(m, keys...), nil
+}
+
+// Omit -
+func (f *CollFuncs) Omit(args ...interface{}) (map[string]interface{}, error) {
+	m, keys, err := pickOmitArgs(args...)
+	if err != nil {
+		return nil, err
+	}
+	return coll.Omit(m, keys...), nil
+}
diff --git a/funcs/coll_test.go b/funcs/coll_test.go
index 135859ba..da554a13 100644
--- a/funcs/coll_test.go
+++ b/funcs/coll_test.go
@@ -23,3 +23,97 @@ func TestFlatten(t *testing.T) {
 	assert.NoError(t, err)
 	assert.EqualValues(t, []interface{}{1, []int{2}, 3}, out)
 }
+
+func TestPick(t *testing.T) {
+	c := &CollFuncs{}
+
+	_, err := c.Pick()
+	assert.Error(t, err)
+
+	_, err = c.Pick("")
+	assert.Error(t, err)
+
+	_, err = c.Pick("foo", nil)
+	assert.Error(t, err)
+
+	_, err = c.Pick("foo", "bar")
+	assert.Error(t, err)
+
+	_, err = c.Pick(map[string]interface{}{}, "foo", "bar", map[string]interface{}{})
+	assert.Error(t, err)
+
+	in := map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+		"":    "baz",
+	}
+	out, err := c.Pick("baz", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, map[string]interface{}{}, out)
+
+	expected := map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+	}
+	out, err = c.Pick("foo", "bar", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, expected, out)
+
+	expected = map[string]interface{}{
+		"": "baz",
+	}
+	out, err = c.Pick("", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, expected, out)
+
+	out, err = c.Pick("foo", "bar", "", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, in, out)
+}
+
+func TestOmit(t *testing.T) {
+	c := &CollFuncs{}
+
+	_, err := c.Omit()
+	assert.Error(t, err)
+
+	_, err = c.Omit("")
+	assert.Error(t, err)
+
+	_, err = c.Omit("foo", nil)
+	assert.Error(t, err)
+
+	_, err = c.Omit("foo", "bar")
+	assert.Error(t, err)
+
+	_, err = c.Omit(map[string]interface{}{}, "foo", "bar", map[string]interface{}{})
+	assert.Error(t, err)
+
+	in := map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+		"":    "baz",
+	}
+	out, err := c.Omit("baz", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, in, out)
+
+	expected := map[string]interface{}{
+		"foo": "bar",
+		"bar": true,
+	}
+	out, err = c.Omit("", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, expected, out)
+
+	expected = map[string]interface{}{
+		"": "baz",
+	}
+	out, err = c.Omit("foo", "bar", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, expected, out)
+
+	out, err = c.Omit("foo", "bar", "", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, map[string]interface{}{}, out)
+}
diff --git a/tests/integration/collection_test.go b/tests/integration/collection_test.go
index f593777a..f21adcc1 100644
--- a/tests/integration/collection_test.go
+++ b/tests/integration/collection_test.go
@@ -118,3 +118,15 @@ func (s *CollSuite) TestFlatten(c *C) {
 		"-i", "{{ `"+in+"` | jsonArray | coll.Flatten 2 | toJSON }}"))
 	result.Assert(c, icmd.Expected{ExitCode: 0, Out: "[1,2,3,4,[[5],6],7]"})
 }
+
+func (s *CollSuite) TestPick(c *C) {
+	result := icmd.RunCmd(icmd.Command(GomplateBin,
+		"-i", `{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}{{ coll.Pick "foo" "baz" $data }}`))
+	result.Assert(c, icmd.Expected{ExitCode: 0, Out: "map[baz:3 foo:1]"})
+}
+
+func (s *CollSuite) TestOmit(c *C) {
+	result := icmd.RunCmd(icmd.Command(GomplateBin,
+		"-i", `{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}{{ coll.Omit "foo" "baz" $data }}`))
+	result.Assert(c, icmd.Expected{ExitCode: 0, Out: "map[bar:2]"})
+}