diff --git a/map.go b/map.go index 9c0ac482..f1ac54b9 100644 --- a/map.go +++ b/map.go @@ -49,8 +49,8 @@ func PickBy[K comparable, V any](in map[K]V, predicate func(key K, value V) bool // Play: https://go.dev/play/p/R1imbuci9qU func PickByKeys[K comparable, V any](in map[K]V, keys []K) map[K]V { r := map[K]V{} - for k, v := range in { - if Contains(keys, k) { + for _, k := range keys { + if v, ok := in[k]; ok { r[k] = v } } @@ -86,9 +86,10 @@ func OmitBy[K comparable, V any](in map[K]V, predicate func(key K, value V) bool func OmitByKeys[K comparable, V any](in map[K]V, keys []K) map[K]V { r := map[K]V{} for k, v := range in { - if !Contains(keys, k) { - r[k] = v - } + r[k] = v + } + for _, k := range keys { + delete(r, k) } return r } diff --git a/map_test.go b/map_test.go index 419ca9fb..c9154425 100644 --- a/map_test.go +++ b/map_test.go @@ -55,7 +55,7 @@ func TestPickByKeys(t *testing.T) { t.Parallel() is := assert.New(t) - r1 := PickByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz"}) + r1 := PickByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz", "qux"}) is.Equal(r1, map[string]int{"foo": 1, "baz": 3}) } @@ -84,7 +84,7 @@ func TestOmitByKeys(t *testing.T) { t.Parallel() is := assert.New(t) - r1 := OmitByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz"}) + r1 := OmitByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"foo", "baz", "qux"}) is.Equal(r1, map[string]int{"bar": 2}) }