New functions coll.Pick and coll.Omit

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

2 files changed, 133 insertions, 0 deletions

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)
+}