fn: Added map functions

Signed-off-by: Ononiwu Maureen <59079323+Chinwendu20@users.noreply.github.com>
This commit is contained in:
Ononiwu Maureen 2024-05-05 21:05:29 +01:00
parent 399ea864da
commit f88f120e91
No known key found for this signature in database
GPG key ID: 82470E85193AD2DA
2 changed files with 179 additions and 0 deletions

44
fn/map.go Normal file
View file

@ -0,0 +1,44 @@
package fn
import (
"fmt"
"golang.org/x/exp/maps"
)
// KeySet converts a map into a Set containing the keys of the map.
func KeySet[K comparable, V any](m map[K]V) Set[K] {
return NewSet(maps.Keys(m)...)
}
// NewSubMapIntersect returns a sub-map of `m` containing only the keys found in
// both `m` and the `keys` slice.
func NewSubMapIntersect[K comparable, V any](m map[K]V, keys []K) map[K]V {
result := make(map[K]V)
for _, k := range keys {
v, ok := m[k]
if !ok {
continue
}
result[k] = v
}
return result
}
// NewSubMap creates a sub-map from a given map using specified keys. It errors
// if any of the keys is not found in the map.
func NewSubMap[K comparable, V any](m map[K]V, keys []K) (map[K]V, error) {
result := make(map[K]V, len(keys))
for _, k := range keys {
v, ok := m[k]
if !ok {
return nil, fmt.Errorf("NewSubMap: missing key %v", k)
}
result[k] = v
}
return result, nil
}

135
fn/map_test.go Normal file
View file

@ -0,0 +1,135 @@
package fn
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestKeySet(t *testing.T) {
testMap := map[string]int{"a": 1, "b": 2, "c": 3}
expected := NewSet([]string{"a", "b", "c"}...)
require.Equal(t, expected, KeySet(testMap))
}
// TestNewSubMap tests the NewSubMap function with various input cases.
func TestNewSubMap(t *testing.T) {
tests := []struct {
name string
original map[int]string
keys []int
expected map[int]string
wantErr bool
}{
{
name: "Successful submap creation",
original: map[int]string{
1: "apple",
2: "banana",
3: "cherry",
},
keys: []int{1, 3},
expected: map[int]string{
1: "apple",
3: "cherry",
},
wantErr: false,
},
{
name: "Key not found",
original: map[int]string{
1: "apple",
2: "banana",
},
keys: []int{1, 4},
expected: nil,
wantErr: true,
},
{
name: "Empty keys list",
original: map[int]string{
1: "apple",
2: "banana",
},
keys: []int{},
expected: map[int]string{},
wantErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result, err := NewSubMap(tt.original, tt.keys)
if tt.wantErr {
require.ErrorContains(
t, err, "NewSubMap: missing key",
)
require.Nil(t, result)
return
}
require.NoError(t, err)
require.Equal(t, tt.expected, result)
})
}
}
// TestNewSubMapIntersect tests the NewSubMapIntersect function for correctness.
func TestNewSubMapIntersect(t *testing.T) {
tests := []struct {
name string
original map[int]string
keys []int
expected map[int]string
}{
{
name: "Successful intersection",
original: map[int]string{
1: "apple",
2: "banana",
3: "cherry",
4: "date",
},
keys: []int{2, 3, 5},
expected: map[int]string{
2: "banana",
3: "cherry",
},
},
{
name: "No intersection",
original: map[int]string{
1: "apple",
2: "banana",
},
keys: []int{3, 4},
expected: map[int]string{},
},
{
name: "Empty original map",
original: map[int]string{},
keys: []int{1, 2},
expected: map[int]string{},
},
{
name: "Empty keys list",
original: map[int]string{
1: "apple",
2: "banana",
},
keys: []int{},
expected: map[int]string{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(
t, tt.expected,
NewSubMapIntersect(tt.original, tt.keys))
})
}
}