Skip to content

Commit

Permalink
Merge pull request #8729 from Chinwendu20/fn
Browse files Browse the repository at this point in the history
Add new  functions to the fn package
  • Loading branch information
guggero committed May 10, 2024
2 parents 4a9ab6e + 30c9b86 commit 4256260
Show file tree
Hide file tree
Showing 4 changed files with 314 additions and 0 deletions.
44 changes: 44 additions & 0 deletions 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
}
135 changes: 135 additions & 0 deletions 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))
})
}
}
30 changes: 30 additions & 0 deletions fn/slice.go
@@ -1,5 +1,13 @@
package fn

import "golang.org/x/exp/constraints"

// Number is a type constraint for all numeric types in Go (integers,
// float and complex numbers)
type Number interface {
constraints.Integer | constraints.Float | constraints.Complex
}

// All returns true when the supplied predicate evaluates to true for all of
// the values in the slice.
func All[A any](pred func(A) bool, s []A) bool {
Expand Down Expand Up @@ -168,3 +176,25 @@ func ZipWith[A, B, C any](f func(A, B) C, a []A, b []B) []C {

return res
}

// SliceToMap converts a slice to a map using the provided key and value
// functions.
func SliceToMap[A any, K comparable, V any](s []A, keyFunc func(A) K,
valueFunc func(A) V) map[K]V {

res := make(map[K]V, len(s))
for _, val := range s {
key := keyFunc(val)
value := valueFunc(val)
res[key] = value
}

return res
}

// Sum calculates the sum of a slice of numbers, `items`.
func Sum[B Number](items []B) B {
return Foldl(func(a, b B) B {
return a + b
}, 0, items)
}
105 changes: 105 additions & 0 deletions fn/slice_test.go
@@ -1,6 +1,7 @@
package fn

import (
"fmt"
"slices"
"testing"

Expand Down Expand Up @@ -136,3 +137,107 @@ func TestZipWith(t *testing.T) {
z, []bool{false, true, false, false, false},
))
}

// TestSum checks if the Sum function correctly calculates the sum of the
// numbers in the slice.
func TestSum(t *testing.T) {
tests := []struct {
name string
items interface{}
result interface{}
}{
{
name: "Sum of positive integers",
items: []int{1, 2, 3},
result: 6,
},
{
name: "Sum of negative integers",
items: []int{-1, -2, -3},
result: -6,
},
{
name: "Sum of float numbers",
items: []float64{1.1, 2.2, 3.3},
result: 6.6,
},
{
name: "Sum of complex numbers",
items: []complex128{
complex(1, 1),
complex(2, 2),
complex(3, 3),
},
result: complex(6, 6),
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
switch v := tt.items.(type) {
case []int:
require.Equal(t, tt.result, Sum(v))
case []float64:
require.Equal(t, tt.result, Sum(v))
case []complex128:
require.Equal(t, tt.result, Sum(v))
}
})
}
}

// TestSliceToMap tests the SliceToMap function.
func TestSliceToMap(t *testing.T) {
tests := []struct {
name string
slice []int
keyFunc func(int) int
valueFunc func(int) string
expected map[int]string
}{
{
name: "Integers to string map",
slice: []int{1, 2, 3},
keyFunc: func(a int) int { return a },
valueFunc: func(a int) string {
return fmt.Sprintf("Value%d", a)
},
expected: map[int]string{
1: "Value1",
2: "Value2",
3: "Value3",
},
},
{
name: "Duplicates in slice",
slice: []int{1, 2, 2, 3},
keyFunc: func(a int) int { return a },
valueFunc: func(a int) string {
return fmt.Sprintf("Value%d", a)
},
expected: map[int]string{
1: "Value1",
2: "Value2",
3: "Value3",
},
},
{
name: "Empty slice",
slice: []int{},
keyFunc: func(a int) int { return a },
valueFunc: func(a int) string {
return fmt.Sprintf("Value%d", a)
},
expected: map[int]string{},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(
t, tt.expected,
SliceToMap(tt.slice, tt.keyFunc, tt.valueFunc),
)
})
}
}

0 comments on commit 4256260

Please sign in to comment.