diff --git a/fn/map.go b/fn/map.go new file mode 100644 index 000000000..071da9872 --- /dev/null +++ b/fn/map.go @@ -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 +} diff --git a/fn/map_test.go b/fn/map_test.go new file mode 100644 index 000000000..b0694bde6 --- /dev/null +++ b/fn/map_test.go @@ -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)) + }) + } +}