From b570a1a539fe0e2c8ac2542639000b6f0fe5c735 Mon Sep 17 00:00:00 2001
From: Dave Henderson <dhenderson@gmail.com>
Date: Sun, 12 May 2024 12:27:44 -0400
Subject: [PATCH] feat: coll.Pick and coll.Omit now support slice of strings as
 input

Signed-off-by: Dave Henderson <dhenderson@gmail.com>
---
 docs-src/content/functions/coll.yml           | 18 +++++++--
 internal/funcs/coll.go                        | 13 +++++++
 internal/funcs/coll_test.go                   | 39 +++++++++++++++++++
 internal/tests/integration/collection_test.go |  2 +
 4 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/docs-src/content/functions/coll.yml b/docs-src/content/functions/coll.yml
index 4302239a7..060271308 100644
--- a/docs-src/content/functions/coll.yml
+++ b/docs-src/content/functions/coll.yml
@@ -410,7 +410,7 @@ funcs:
     description: |
       Given a map, returns a new map with any entries that have the given keys.
 
-      All keys are converted to strings.
+      The keys can either be separate arguments, or a slice (since v4.0.0).
 
       This is the inverse of [`coll.Omit`](#coll-omit).
 
@@ -419,7 +419,7 @@ funcs:
     arguments:
       - name: keys...
         required: true
-        description: the keys to match
+        description: the keys (strings) to match
       - name: map
         required: true
         description: the map to pick from
@@ -428,12 +428,17 @@ funcs:
         $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
         {{ coll.Pick "foo" "baz" $data }}'
         map[baz:3 foo:1]
+      - |
+        $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+        {{ $keys := coll.Slice "foo" "baz" }}
+        {{ coll.Pick $keys $data }}'
+        map[baz:3 foo:1]
   - name: coll.Omit
     released: v3.7.0
     description: |
       Given a map, returns a new map without any entries that have the given keys.
 
-      All keys are converted to strings.
+      The keys can either be separate arguments, or a slice (since v4.0.0).
 
       This is the inverse of [`coll.Pick`](#coll-pick).
 
@@ -442,7 +447,7 @@ funcs:
     arguments:
       - name: keys...
         required: true
-        description: the keys to match
+        description: the keys (strings) to match
       - name: map
         required: true
         description: the map to omit from
@@ -451,3 +456,8 @@ funcs:
         $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
         {{ coll.Omit "foo" "baz" $data }}'
         map[bar:2]
+      - |
+        $ gomplate -i '{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}
+        {{ $keys := coll.Slice "foo" "baz" }}
+        {{ coll.Omit $keys $data }}'
+        map[bar:2]
diff --git a/internal/funcs/coll.go b/internal/funcs/coll.go
index a9431f9d6..9e08aeacf 100644
--- a/internal/funcs/coll.go
+++ b/internal/funcs/coll.go
@@ -172,6 +172,19 @@ func pickOmitArgs(args ...interface{}) (map[string]interface{}, []string, error)
 		return nil, nil, fmt.Errorf("wrong map type: must be map[string]interface{}, got %T", args[len(args)-1])
 	}
 
+	// special-case - if there's only one key and it's a slice, expand it
+	if len(args) == 2 {
+		if reflect.TypeOf(args[0]).Kind() == reflect.Slice {
+			sl := reflect.ValueOf(args[0])
+			expandedArgs := make([]interface{}, sl.Len()+1)
+			for i := 0; i < sl.Len(); i++ {
+				expandedArgs[i] = sl.Index(i).Interface()
+			}
+			expandedArgs[len(expandedArgs)-1] = m
+			args = expandedArgs
+		}
+	}
+
 	keys := make([]string, len(args)-1)
 	for i, v := range args[0 : len(args)-1] {
 		k, ok := v.(string)
diff --git a/internal/funcs/coll_test.go b/internal/funcs/coll_test.go
index 1a0d68597..fe0278d76 100644
--- a/internal/funcs/coll_test.go
+++ b/internal/funcs/coll_test.go
@@ -94,6 +94,19 @@ func TestPick(t *testing.T) {
 	out, err = c.Pick("foo", "bar", "", in)
 	require.NoError(t, err)
 	assert.EqualValues(t, in, out)
+
+	t.Run("supports slice key", func(t *testing.T) {
+		t.Parallel()
+
+		in := map[string]interface{}{
+			"foo": "bar",
+			"bar": true,
+			"":    "baz",
+		}
+		out, err := c.Pick([]string{"foo", "bar"}, in)
+		require.NoError(t, err)
+		assert.EqualValues(t, map[string]interface{}{"foo": "bar", "bar": true}, out)
+	})
 }
 
 func TestOmit(t *testing.T) {
@@ -143,6 +156,32 @@ func TestOmit(t *testing.T) {
 	out, err = c.Omit("foo", "bar", "", in)
 	require.NoError(t, err)
 	assert.EqualValues(t, map[string]interface{}{}, out)
+
+	t.Run("supports slice of strings", func(t *testing.T) {
+		t.Parallel()
+
+		in := map[string]interface{}{
+			"foo": "bar",
+			"bar": true,
+			"":    "baz",
+		}
+		out, err := c.Omit([]string{"foo", "bar"}, in)
+		require.NoError(t, err)
+		assert.EqualValues(t, map[string]interface{}{"": "baz"}, out)
+	})
+
+	t.Run("supports slice of interface{}", func(t *testing.T) {
+		t.Parallel()
+
+		in := map[string]interface{}{
+			"foo": "bar",
+			"bar": true,
+			"":    "baz",
+		}
+		out, err := c.Omit([]interface{}{"foo", "bar"}, in)
+		require.NoError(t, err)
+		assert.EqualValues(t, map[string]interface{}{"": "baz"}, out)
+	})
 }
 
 func TestGoSlice(t *testing.T) {
diff --git a/internal/tests/integration/collection_test.go b/internal/tests/integration/collection_test.go
index 5175ef64f..45c33ad13 100644
--- a/internal/tests/integration/collection_test.go
+++ b/internal/tests/integration/collection_test.go
@@ -93,10 +93,12 @@ func TestColl_Flatten(t *testing.T) {
 
 func TestColl_Pick(t *testing.T) {
 	inOutTest(t, `{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}{{ coll.Pick "foo" "baz" $data }}`, "map[baz:3 foo:1]")
+	inOutTest(t, `{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}{{ coll.Pick (coll.Slice "foo" "baz") $data }}`, "map[baz:3 foo:1]")
 }
 
 func TestColl_Omit(t *testing.T) {
 	inOutTest(t, `{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}{{ coll.Omit "foo" "baz" $data }}`, "map[bar:2]")
+	inOutTest(t, `{{ $data := dict "foo" 1 "bar" 2 "baz" 3 }}{{ coll.Omit (coll.Slice "foo" "baz") $data }}`, "map[bar:2]")
 }
 
 func TestColl_JQ(t *testing.T) {