New functions coll.Pick and coll.Omit

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

2 files changed, 79 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", ""))
+}