diff --git a/README.md b/README.md index 73c33223..d2ddf9a8 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,7 @@ Supported helpers for maps: - [PickBy](#pickby) - [PickByKeys](#pickbykeys) - [PickByValues](#pickbyvalues) +- [ValuesByKeys](#ValuesByKeys) - [OmitBy](#omitby) - [OmitByKeys](#omitbykeys) - [OmitByValues](#omitbyvalues) @@ -1173,6 +1174,18 @@ m := lo.PickByValues(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []int{1, 3}) [[play](https://go.dev/play/p/1zdzSvbfsJc)] +### ValuesByKeys + +Returns an array of values in the same order as the given keys. +If the key in keys slice but not in the map, return error +```go + +m, err := lo.ValuesByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"baz", "foo", "bar"}) +// []int{3,1,2} +``` + +[[play](https://go.dev/play/p/PmmPDo1AqWl)] + ### OmitBy Returns same map type filtered by given predicate. diff --git a/map.go b/map.go index d8feb434..2cd60c36 100644 --- a/map.go +++ b/map.go @@ -1,5 +1,9 @@ package lo +import ( + "fmt" +) + // Keys creates an array of the map keys. // Play: https://go.dev/play/p/Uu11fHASqrU func Keys[K comparable, V any](in ...map[K]V) []K { @@ -137,6 +141,21 @@ func PickByValues[K comparable, V comparable, Map ~map[K]V](in Map, values []V) return r } +// ValuesByKeys returns an array of values in the same order as the given keys. +// if the key in keys slice but not in the map, return error +// Play: https://go.dev/play/p/J0SvRMzPJo_s +func ValuesByKeys[K comparable, V any](in map[K]V, keys []K) ([]V, error) { + out := make([]V, 0, len(keys)) + for i := range keys { + v, ok := in[keys[i]] + if !ok { + return nil, fmt.Errorf("ValuesByKeys: %v is not in the map", keys[i]) + } + out = append(out, v) + } + return out, nil +} + // OmitBy returns same map type filtered by given predicate. // Play: https://go.dev/play/p/EtBsR43bdsd func OmitBy[K comparable, V any, Map ~map[K]V](in Map, predicate func(key K, value V) bool) Map { diff --git a/map_example_test.go b/map_example_test.go index e9347a04..56149c24 100644 --- a/map_example_test.go +++ b/map_example_test.go @@ -88,6 +88,17 @@ func ExamplePickByValues() { // Output: 2 1 3 } +func ExampleValuesByKeys() { + kv := map[string]int{"foo": 1, "bar": 2, "baz": 3} + + result, err := ValuesByKeys(kv, []string{"baz", "foo", "bar"}) + if err != nil { + return + } + fmt.Printf("%v %v %v", len(result), result[0], result[1]) + // Output: 3 3 1 +} + func ExampleOmitBy() { kv := map[string]int{"foo": 1, "bar": 2, "baz": 3} diff --git a/map_test.go b/map_test.go index fb02bf9a..c2900fb5 100644 --- a/map_test.go +++ b/map_test.go @@ -179,6 +179,24 @@ func TestPickByValues(t *testing.T) { is.IsType(after, before, "type preserved") } +func TestValuesByKeys(t *testing.T) { + t.Parallel() + is := assert.New(t) + + result1, err1 := ValuesByKeys(map[string]int{"foo": 1, "bar": 2, "baz": 3}, []string{"baz", "foo", "bar"}) + is.NoError(err1) + is.Equal(len(result1), 3) + is.ElementsMatch(result1, []int{3, 1, 2}) + + result2, err2 := ValuesByKeys(map[string]int{"": 0, "foobar": 6, "baz": 3}, []string{"baz", "foobar"}) + is.NoError(err2) + is.Equal(len(result2), 2) + is.ElementsMatch(result2, []int{3, 6}) + + _, err3 := ValuesByKeys(map[string]int{"foo": 1, "baz": 3}, []string{"baz", "foo", "bar"}) + is.Errorf(err3, "ValuesByKeys: bar is not in the map") +} + func TestOmitBy(t *testing.T) { t.Parallel() is := assert.New(t)